diff --git a/monai/data/utils.py b/monai/data/utils.py index cf426d4c3e..c424d53e29 100644 --- a/monai/data/utils.py +++ b/monai/data/utils.py @@ -670,6 +670,11 @@ def set_rnd(obj, seed: int) -> int: obj: object to set seed or random state for. seed: set the random state with an integer seed. """ + if isinstance(obj, (tuple, list)): # ZipDataset.data is a list + _seed = seed + for item in obj: + _seed = set_rnd(item, seed=seed) + return seed if _seed == seed else seed + 1 # return a different seed if there are randomizable items if not hasattr(obj, "__dict__"): return seed # no attribute if hasattr(obj, "set_random_state"): diff --git a/tests/test_dataloader.py b/tests/test_dataloader.py index 79126b2dbb..88fd585a32 100644 --- a/tests/test_dataloader.py +++ b/tests/test_dataloader.py @@ -16,9 +16,10 @@ import torch from parameterized import parameterized -from monai.data import CacheDataset, DataLoader, Dataset +from monai.data import CacheDataset, DataLoader, Dataset, ZipDataset from monai.transforms import Compose, DataStatsd, Randomizable, SimulateDelayd -from monai.utils import set_determinism +from monai.utils import convert_to_numpy, set_determinism +from tests.utils import assert_allclose TEST_CASE_1 = [[{"image": np.asarray([1, 2, 3])}, {"image": np.asarray([4, 5])}]] @@ -83,6 +84,15 @@ def test_randomize(self): output.extend(batch.data.numpy().flatten().tolist()) self.assertListEqual(output, [594, 170, 524, 778, 370, 906, 292, 589, 762, 763, 156, 886, 42, 405, 221, 166]) + def test_zipdataset(self): + dataset = ZipDataset([_RandomDataset(), ZipDataset([_RandomDataset(), _RandomDataset()])]) + dataloader = DataLoader(dataset, batch_size=2, num_workers=2) + output = [] + for _ in range(2): + for batch in dataloader: + output.extend([convert_to_numpy(batch, wrap_sequence=False)]) + assert_allclose(np.stack(output).flatten()[:7], np.array([594, 170, 594, 170, 594, 170, 524])) + if __name__ == "__main__": unittest.main()