Add option to specify number of steps in UniformSampler (#550)

Summary:
Pull Request resolved: #550

Add option to specify number of steps in UniformSampler

Differential Revision: D42034268

fbshipit-source-id: c0c36b8f3bfe8903b717e2b2203ae7b68f493074

opacus/utils/uniform_sampler.py

@@ -25,12 +25,15 @@ class UniformWithReplacementSampler(Sampler[List[int]]):
     Each sample is selected with a probability equal to ``sample_rate``.
     """
 
-    def __init__(self, *, num_samples: int, sample_rate: float, generator=None):
+    def __init__(
+        self, *, num_samples: int, sample_rate: float, generator=None, steps=None
+    ):
         r"""
         Args:
             num_samples: number of samples to draw.
             sample_rate: probability used in sampling.
             generator: Generator used in sampling.
+            steps: Number of steps (iterations of the Sampler)
         """
         self.num_samples = num_samples
         self.sample_rate = sample_rate
@@ -42,11 +45,16 @@ def __init__(self, *, num_samples: int, sample_rate: float, generator=None):
                 "value, but got num_samples={}".format(self.num_samples)
             )
 
+        if steps is not None:
+            self.steps = steps
+        else:
+            self.steps = int(1 / self.sample_rate)
+
     def __len__(self):
-        return int(1 / self.sample_rate)
+        return self.steps
 
     def __iter__(self):
-        num_batches = int(1 / self.sample_rate)
+        num_batches = self.steps
         while num_batches > 0:
             mask = (
                 torch.rand(self.num_samples, generator=self.generator)
@@ -82,6 +90,7 @@ def __init__(
         sample_rate: float,
         shuffle: bool = True,
         shuffle_seed: int = 0,
+        steps: int = None,
         generator=None,
     ):
         """
@@ -117,7 +126,10 @@ def __init__(
             self.num_samples += 1
 
         # Number of batches: same as non-distributed Poisson sampling, but each batch is smaller
-        self.num_batches = int(1 / self.sample_rate)
+        if steps is not None:
+            self.num_batches = steps
+        else:
+            self.num_batches = int(1 / self.sample_rate)
 
     def __iter__(self):
         if self.shuffle: