Patch for pytorch#40026 RandomSampler generates samples one at a time when replacement=True (pytorch#41682)

Summary:
Fix pytorch#32530
Fix/Patch pytorch#40026

Resubmit this patch and fix the type error.

Force the input type to `manual_seed()` in `sampler.py` to be `int`.

ezyang

Pull Request resolved: pytorch#41682

Reviewed By: izdeby

Differential Revision: D22665477

Pulled By: ezyang

fbshipit-source-id: 1725c8aa742c31e74321f20448f4b6a392afb38d

Author: ydaiming

test/test_dataloader.py

@@ -1288,22 +1288,24 @@ def test_random_sampler(self):
         def sample_stat(sampler, num_samples):
             counts = Counter(sampler)
             count_repeated = sum(val > 1 for val in counts.values())
-            return (count_repeated, min(counts.keys()), max(counts.keys()))
+            return (count_repeated, min(counts.keys()), max(counts.keys()), sum(counts.values()))
 
         # test sample with replacement
         n = len(self.dataset) + 1  # ensure at least one sample is drawn more than once
         sampler_with_replacement = RandomSampler(self.dataset, replacement=True, num_samples=n)
-        count_repeated, minval, maxval = sample_stat(sampler_with_replacement, n)
+        count_repeated, minval, maxval, count_total = sample_stat(sampler_with_replacement, n)
         self.assertTrue(count_repeated > 0)
         self.assertTrue(minval >= 0)
         self.assertTrue(maxval < len(self.dataset))
+        self.assertTrue(count_total == n)
 
         # test sample without replacement
         sampler_without_replacement = RandomSampler(self.dataset)
-        count_repeated, minval, maxval = sample_stat(sampler_without_replacement, len(self.dataset))
+        count_repeated, minval, maxval, count_total = sample_stat(sampler_without_replacement, len(self.dataset))
         self.assertTrue(count_repeated == 0)
         self.assertTrue(minval == 0)
         self.assertTrue(maxval == len(self.dataset) - 1)
+        self.assertTrue(count_total == len(self.dataset))
 
         # raise error when replacement=False and num_samples is not None
         self.assertRaises(ValueError, lambda: RandomSampler(self.dataset, num_samples=len(self.dataset)))

torch/utils/data/sampler.py

@@ -76,7 +76,7 @@ class RandomSampler(Sampler[int]):
 
     Arguments:
         data_source (Dataset): dataset to sample from
-        replacement (bool): samples are drawn with replacement if ``True``, default=``False``
+        replacement (bool): samples are drawn on-demand with replacement if ``True``, default=``False``
         num_samples (int): number of samples to draw, default=`len(dataset)`. This argument
             is supposed to be specified only when `replacement` is ``True``.
         generator (Generator): Generator used in sampling.
@@ -112,10 +112,17 @@ def num_samples(self) -> int:
 
     def __iter__(self):
         n = len(self.data_source)
+        if self.generator is None:
+            generator = torch.Generator()
+            generator.manual_seed(int(torch.empty((), dtype=torch.int64).random_().item()))
+        else:
+            generator = self.generator
         if self.replacement:
-            rand_tensor = torch.randint(high=n, size=(self.num_samples,), dtype=torch.int64, generator=self.generator)
-            return iter(rand_tensor.tolist())
-        return iter(torch.randperm(n, generator=self.generator).tolist())
+            for _ in range(self.num_samples // 32):
+                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()
 
     def __len__(self):
         return self.num_samples