4569 consistent rng for data loader (#5067)

Signed-off-by: Wenqi Li <[email protected]>

Fixes #4569 

### Description
with different `num_workers` the batch sampler's rng should be consistent, this PR keeps the initial random seed to achieve it.

### Status
**Ready**

### Types of changes
<!--- Put an `x` in all the boxes that apply, and remove the not applicable items -->
- [x] Non-breaking change (fix or new feature that would not break existing functionality).
- [ ] Breaking change (fix or new feature that would cause existing functionality to change). # the previous results for num_wokers==0 might change
- [x] New tests added to cover the changes.
- [ ] Integration tests passed locally by running `./runtests.sh -f -u --net --coverage`.
- [ ] Quick tests passed locally by running `./runtests.sh --quick --unittests  --disttests`.
- [x] In-line docstrings updated.
- [x] Documentation updated, tested `make html` command in the `docs/` folder.

Signed-off-by: Wenqi Li <[email protected]>

Author: wyli

monai/data/dataloader.py

@@ -66,6 +66,8 @@ def __len__(self):
         num_workers: how many subprocesses to use for data
             loading. ``0`` means that the data will be loaded in the main process.
             (default: ``0``)
+        collate_fn: default to :py:func:`monai.data.utils.list_data_collate`.
+        worker_init_fn: default to :py:func:`monai.data.utils.worker_init_fn`.
         kwargs: other parameters for PyTorch DataLoader.
     """
 
@@ -74,11 +76,14 @@ def __init__(self, dataset: Dataset, num_workers: int = 0, **kwargs) -> None:
             # when num_workers > 0, random states are determined by worker_init_fn
             # this is to make the behavior consistent when num_workers == 0
             # torch.int64 doesn't work well on some versions of windows
-            _seed = torch.empty((), dtype=torch.int32).random_(generator=None).item()
+            _g = torch.random.default_generator if kwargs.get("generator", None) is None else kwargs["generator"]
+            init_seed = _g.initial_seed()
+            _seed = torch.empty((), dtype=torch.int64).random_(generator=_g).item()
             set_rnd(dataset, int(_seed))
+            _g.manual_seed(init_seed)
         if "collate_fn" not in kwargs:
-            kwargs.update({"collate_fn": list_data_collate})
+            kwargs["collate_fn"] = list_data_collate
         if "worker_init_fn" not in kwargs:
-            kwargs.update({"worker_init_fn": worker_init_fn})
+            kwargs["worker_init_fn"] = worker_init_fn
 
         super().__init__(dataset=dataset, num_workers=num_workers, **kwargs)

tests/test_compose.py

@@ -169,16 +169,16 @@ def test_data_loader(self):
         set_determinism(seed=123)
         train_loader = DataLoader(train_ds, num_workers=0)
         out_1 = next(iter(train_loader))
-        self.assertAlmostEqual(out_1.cpu().item(), 0.84291356)
+        self.assertAlmostEqual(out_1.cpu().item(), 0.0409280)
 
         if sys.platform != "win32":  # skip multi-worker tests on win32
             train_loader = DataLoader(train_ds, num_workers=1)
             out_1 = next(iter(train_loader))
-            self.assertAlmostEqual(out_1.cpu().item(), 0.180814653)
+            self.assertAlmostEqual(out_1.cpu().item(), 0.78663897075)
 
             train_loader = DataLoader(train_ds, num_workers=2)
             out_1 = next(iter(train_loader))
-            self.assertAlmostEqual(out_1.cpu().item(), 0.04293707)
+            self.assertAlmostEqual(out_1.cpu().item(), 0.785907334)
         set_determinism(None)
 
     def test_data_loader_2(self):
@@ -191,16 +191,16 @@ def test_data_loader_2(self):
 
         train_loader = DataLoader(train_ds, num_workers=0)
         out_2 = next(iter(train_loader))
-        self.assertAlmostEqual(out_2.cpu().item(), 0.7858843729)
+        self.assertAlmostEqual(out_2.cpu().item(), 0.98921915918)
 
         if sys.platform != "win32":  # skip multi-worker tests on win32
             train_loader = DataLoader(train_ds, num_workers=1)
             out_2 = next(iter(train_loader))
-            self.assertAlmostEqual(out_2.cpu().item(), 0.305763411)
+            self.assertAlmostEqual(out_2.cpu().item(), 0.32985207)
 
             train_loader = DataLoader(train_ds, num_workers=2)
             out_1 = next(iter(train_loader))
-            self.assertAlmostEqual(out_1.cpu().item(), 0.131966779)
+            self.assertAlmostEqual(out_1.cpu().item(), 0.28602141572)
         set_determinism(None)
 
     def test_flatten_and_len(self):

tests/test_dataloader.py

@@ -75,14 +75,17 @@ def setUp(self):
     def tearDown(self):
         set_determinism(None)
 
-    def test_randomize(self):
+    @parameterized.expand([[1], [0]])
+    def test_randomize(self, workers):
+        set_determinism(0)
         dataset = _RandomDataset()
-        dataloader = DataLoader(dataset, batch_size=2, num_workers=3)
+        dataloader = DataLoader(dataset, batch_size=2, shuffle=True, num_workers=workers)
         output = []
-        for _ in range(2):
+        for _ in range(1):  # need persistent workers for reproducibility of num_workers 0, 1
             for batch in dataloader:
                 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])
+        set_determinism(None)
+        self.assertListEqual(output, [594, 170, 292, 589, 153, 811, 21, 550])
 
     def test_zipdataset(self):
         dataset = ZipDataset([_RandomDataset(), ZipDataset([_RandomDataset(), _RandomDataset()])])

tests/test_grid_dataset.py

@@ -60,7 +60,7 @@ def test_loading_array(self):
             np.testing.assert_equal(tuple(item[0].shape), (2, 1, 2, 2))
         np.testing.assert_allclose(
             item[0],
-            np.array([[[[8.0577, 9.0577], [12.0577, 13.0577]]], [[[10.5540, 11.5540], [14.5540, 15.5540]]]]),
+            np.array([[[[8.240326, 9.240326], [12.240326, 13.240326]]], [[[10.1624, 11.1624], [14.1624, 15.1624]]]]),
             rtol=1e-4,
         )
         np.testing.assert_allclose(item[1], np.array([[[0, 1], [2, 4], [0, 2]], [[0, 1], [2, 4], [2, 4]]]), rtol=1e-5)
@@ -69,7 +69,9 @@ def test_loading_array(self):
                 np.testing.assert_equal(tuple(item[0].shape), (2, 1, 2, 2))
             np.testing.assert_allclose(
                 item[0],
-                np.array([[[[7.6533, 8.6533], [11.6533, 12.6533]]], [[[9.8524, 10.8524], [13.8524, 14.8524]]]]),
+                np.array(
+                    [[[[7.723618, 8.723618], [11.723618, 12.723618]]], [[[10.7175, 11.7175], [14.7175, 15.7175]]]]
+                ),
                 rtol=1e-3,
             )
             np.testing.assert_allclose(
@@ -102,7 +104,7 @@ def test_loading_dict(self):
             self.assertListEqual(item[0]["metadata"], ["test string", "test string"])
         np.testing.assert_allclose(
             item[0]["image"],
-            np.array([[[[8.0577, 9.0577], [12.0577, 13.0577]]], [[[10.5540, 11.5540], [14.5540, 15.5540]]]]),
+            np.array([[[[8.240326, 9.240326], [12.240326, 13.240326]]], [[[10.1624, 11.1624], [14.1624, 15.1624]]]]),
             rtol=1e-4,
         )
         np.testing.assert_allclose(item[1], np.array([[[0, 1], [2, 4], [0, 2]], [[0, 1], [2, 4], [2, 4]]]), rtol=1e-5)
@@ -111,7 +113,9 @@ def test_loading_dict(self):
                 np.testing.assert_equal(item[0]["image"].shape, (2, 1, 2, 2))
             np.testing.assert_allclose(
                 item[0]["image"],
-                np.array([[[[7.6533, 8.6533], [11.6533, 12.6533]]], [[[9.8524, 10.8524], [13.8524, 14.8524]]]]),
+                np.array(
+                    [[[[7.723618, 8.723618], [11.723618, 12.723618]]], [[[10.7175, 11.7175], [14.7175, 15.7175]]]]
+                ),
                 rtol=1e-3,
             )
             np.testing.assert_allclose(

tests/test_patch_dataset.py

@@ -59,7 +59,7 @@ def test_loading_array(self):
         np.testing.assert_allclose(
             item[0],
             np.array(
-                [[[1.338681, 2.338681, 3.338681], [5.338681, 6.338681, 7.338681], [9.338681, 10.338681, 11.338681]]]
+                [[[-0.593095, 0.406905, 1.406905], [3.406905, 4.406905, 5.406905], [7.406905, 8.406905, 9.406905]]]
             ),
             rtol=1e-5,
         )
@@ -69,13 +69,7 @@ def test_loading_array(self):
             np.testing.assert_allclose(
                 item[0],
                 np.array(
-                    [
-                        [
-                            [4.957847, 5.957847, 6.957847],
-                            [8.957847, 9.957847, 10.957847],
-                            [12.957847, 13.957847, 14.957847],
-                        ]
-                    ]
+                    [[[0.234308, 1.234308, 2.234308], [4.234308, 5.234308, 6.234308], [8.234308, 9.234308, 10.234308]]]
                 ),
                 rtol=1e-5,
             )

tests/test_thread_buffer.py

@@ -13,9 +13,12 @@
 import time
 import unittest
 
+import torch
+
 from monai.data import DataLoader, Dataset, ThreadBuffer, ThreadDataLoader
 from monai.transforms import Compose, SimulateDelayd
-from monai.utils import PerfContext
+from monai.utils import PerfContext, set_determinism
+from tests.utils import assert_allclose
 
 
 class TestDataLoader(unittest.TestCase):
@@ -53,6 +56,19 @@ def test_dataloader(self):
             self.assertEqual(d["label"][0], "spleen_label_19.nii.gz")
             self.assertEqual(d["label"][1], "spleen_label_31.nii.gz")
 
+    def test_deterministic(self):
+        set_determinism(0)
+        res_1 = list(ThreadDataLoader(torch.arange(5), batch_size=2, buffer_size=2, shuffle=True, num_workers=0))
+
+        set_determinism(0)
+        num_workers = 2 if sys.platform == "linux" else 1
+        res_2 = list(
+            ThreadDataLoader(torch.arange(5), batch_size=2, buffer_size=3, shuffle=True, num_workers=num_workers)
+        )
+
+        set_determinism(None)
+        assert_allclose(torch.cat(res_1), torch.cat(res_2), type_test=False)
+
     def test_time(self):
         dataset = Dataset(data=self.datalist * 2, transform=self.transform)  # contains data for 2 batches
         dataloader = DataLoader(dataset=dataset, batch_size=2, num_workers=0)

tests/testing_data/transform_metatensor_cases.yaml

@@ -180,7 +180,7 @@ TEST_CASE_3:
 
 TEST_CASE_1_answer:
   load_shape: [1, 1, 33, 45, 54]
-  affine: "$np.array([[-2, 0, 0, 34], [0, 2, 0, -64], [0, 0, 2, -54], [0, 0, 0, 1]], dtype=np.float64)"
+  affine: "$np.array([[-2, 0, 0, 30], [0, 2, 0, -62], [0, 0, 2, -48], [0, 0, 0, 1]], dtype=np.float64)"
   inv_affine: "@init_affine"
   inv_shape: "@init_shape"
 
@@ -192,6 +192,6 @@ TEST_CASE_2_answer:
 
 TEST_CASE_3_answer:
   load_shape: [1, 1, 72, 57, 82]
-  affine: "$np.array([[-1.343816, -0.682904, -0.234832, 76.01494], [0.309004, 0.653211, -1.734872, 24.511358], [-0.104049, 1.617199, 0.584171, -56.521294], [0, 0, 0, 1]], dtype=np.float64)"
+  affine: "$np.array([[1.300558,  -0.700765,  -0.511861,  -3.739605], [0.479723,  -1.171149,   1.193079, -50.087933], [0.395736,   1.183532,   0.984201, -80.496605], [0, 0, 0, 1]], dtype=np.float64)"
   inv_affine: "@init_affine"
   inv_shape: "@init_shape"