Fix seeding random module in DataLoader (#7886)

thuyen · soumith · commit 146b951ec5cb · 2018-05-29T15:55:04.000-04:00
* fix seeding random module

* make base seed int

* follow 0.4 idiom

* add a test for random seeding
diff --git a/test/test_utils.py b/test/test_utils.py
@@ -116,6 +116,15 @@ def __len__(self):
         return 10
 
 
+class RandomDatasetMock(object):
+
+    def __getitem__(self, index):
+        return torch.tensor([torch.rand(1).item(), random.uniform(0, 1)])
+
+    def __len__(self):
+        return 1000
+
+
 class TestCheckpoint(TestCase):
 
     # Test whether checkpoint is being triggered or not. For this, we check
@@ -233,6 +242,20 @@ def setUp(self):
         self.dataset = torch.randn(5, 3, 3, 2)
         self.batch_size = 3
 
+    def test_random_seed(self):
+        def run():
+            dataloader = torch.utils.data.DataLoader(RandomDatasetMock(),
+                                                     batch_size=2,
+                                                     num_workers=4,
+                                                     shuffle=True)
+            return next(iter(dataloader))
+
+        torch.manual_seed(2018)
+        x1 = run()
+        torch.manual_seed(2018)
+        x2 = run()
+        self.assertEqual(x1, x2)
+
     def test_single_keep(self):
         dataloader = torch.utils.data.DataLoader(self.dataset,
                                                  batch_size=self.batch_size,
diff --git a/torch/utils/data/dataloader.py b/torch/utils/data/dataloader.py
@@ -246,7 +246,7 @@ def __init__(self, loader):
 
         self.sample_iter = iter(self.batch_sampler)
 
-        base_seed = torch.LongTensor(1).random_()[0]
+        base_seed = torch.LongTensor(1).random_().item()
 
         if self.num_workers > 0:
             self.worker_init_fn = loader.worker_init_fn
