Repository: pytorch/pytorch

Issue:
Sampler should be seeded lazily
## 🐛 Bug

After #63026 is landed, the generator for Sampler is attached to the instance, which helps to serialize the state of Sampler. But, it brings a problem that it will prevent Sampler's generator being seeded before each epoch.

## To Reproduce

Check https://github.com/pytorch/pytorch/pull/63026#issuecomment-902234490

User would expect same result for the sampler without specifying generator input when set seed for each epoch.
```py
sampler = RandomSampler(ds)
torch.manual_seed(0)
l1 = list(sampler)
torch.manual_seed(0)
l2 = list(sampler)
# Expect same
assert l1 == l2
```


cc @SsnL @VitalyFedyunin @ejguan @NivekT


Discussion:

Resolving patch:
diff --git a/torch/utils/data/sampler.py b/torch/utils/data/sampler.py
--- a/torch/utils/data/sampler.py
+++ b/torch/utils/data/sampler.py
@@ -112,15 +112,18 @@ def num_samples(self) -> int:
     def __iter__(self) -> Iterator[int]:
         n = len(self.data_source)
         if self.generator is None:
-            self.generator = torch.Generator()
-            self.generator.manual_seed(int(torch.empty((), dtype=torch.int64).random_().item()))
+            seed = int(torch.empty((), dtype=torch.int64).random_().item())
+            generator = torch.Generator()
+            generator.manual_seed(seed)
+        else:
+            generator = self.generator
 
         if self.replacement:
             for _ in range(self.num_samples // 32):
-                yield from torch.randint(high=n, size=(32,), dtype=torch.int64, generator=self.generator).tolist()
-            yield from torch.randint(high=n, size=(self.num_samples % 32,), dtype=torch.int64, generator=self.generator).tolist()
+                yield from torch.randint(high=n, size=(32,), dtype=torch.int64, generator=generator).tolist()
+            yield from torch.randint(high=n, size=(self.num_samples % 32,), dtype=torch.int64, generator=generator).tolist()
         else:
-            yield from torch.randperm(n, generator=self.generator).tolist()
+            yield from torch.randperm(n, generator=generator).tolist()
 
     def __len__(self) -> int:
         return self.num_samples
@@ -140,7 +143,8 @@ def __init__(self, indices: Sequence[int], generator=None) -> None:
         self.generator = generator
 
     def __iter__(self) -> Iterator[int]:
-        return (self.indices[i] for i in torch.randperm(len(self.indices), generator=self.generator))
+        for i in torch.randperm(len(self.indices), generator=self.generator):
+            yield self.indices[i]
 
     def __len__(self) -> int:
         return len(self.indices)
@@ -183,7 +187,7 @@ def __init__(self, weights: Sequence[float], num_samples: int,
 
     def __iter__(self) -> Iterator[int]:
         rand_tensor = torch.multinomial(self.weights, self.num_samples, self.replacement, generator=self.generator)
-        return iter(rand_tensor.tolist())
+        yield from iter(rand_tensor.tolist())
 
     def __len__(self) -> int:
         return self.num_samples


Test patch:
diff --git a/test/test_dataloader.py b/test/test_dataloader.py
--- a/test/test_dataloader.py
+++ b/test/test_dataloader.py
@@ -1524,6 +1524,28 @@ def test_sampler_reproducibility(self):
         ):
             self.assertEqual(list(fn()), list(fn()))
 
+        for sampler in (
+            RandomSampler(self.dataset, num_samples=5, replacement=True),
+            RandomSampler(self.dataset, replacement=False),
+            WeightedRandomSampler(weights, num_samples=5, replacement=True),
+            WeightedRandomSampler(weights, num_samples=5, replacement=False),
+            SubsetRandomSampler(range(10)),
+        ):
+            torch.manual_seed(0)
+            l1 = list(sampler) + list(sampler)
+
+            torch.manual_seed(0)
+            l2 = list(sampler) + list(sampler)
+            self.assertEqual(l1, l2)
+
+            its = (iter(sampler), iter(sampler))
+            ls = ([], [])
+            for idx in range(len(sampler)):
+                for i in range(2):
+                    if idx == 0:
+                        torch.manual_seed(0)
+                    ls[i].append(next(its[i]))
+            self.assertEqual(ls[0], ls[1])
 
     def _test_sampler(self, **kwargs):
         indices = range(2, 12)  # using a regular iterable