class_docstr: Distributed batch sampler.
  
  Each batch is sampled as follows:
      1. Shuffle the dataset (enabled by default)
      2. Split the dataset among the replicas into chunks of equal size
         (plus or minus one sample)
      3. Each replica selects each sample of its chunk independently
         with probability `sample_rate`
      4. Each replica outputs the selected samples, which form a local batch
  
  The sum of the lengths of the local batches follows a Poisson distribution.
  In particular, the expected length of each local batch is:
  `sample_rate * total_size / num_replicas`
class_docstr_tok: ["Distributed","batch","sampler",".","Each","batch","is","sampled","as","follows",":","1",".","Shuffle","the","dataset","(","enabled","by","default",")","2",".","Split","the","dataset","among","the","replicas","into","chunks","of","equal","size","(","plus","or","minus","one","sample",")","3",".","Each","replica","selects","each","sample","of","its","chunk","independently","with","probability","`","sample_rate","`","4",".","Each","replica","outputs","the","selected","samples",",","which","form","a","local","batch","The","sum","of","the","lengths","of","the","local","batches","follows","a","Poisson","distribution",".","In","particular",",","the","expected","length","of","each","local","batch","is",":","`","sample_rate","*","total_size","\/","num_replicas","`"]
docstr: Args:
      total_size: total number of samples to sample from
      sample_rate: number of samples to draw.
      shuffle: Flag indicating whether apply shuffle when dividing elements
          between workers
      shuffle_seed: Random seed used to shuffle when dividing elements across workers
      generator: torch.Generator() object used as a source of randomness
          when selecting items for the next round on a given worker
docstr_tok: ["Args",":","total_size",":","total","number","of","samples","to","sample","from","sample_rate",":","number","of","samples","to","draw",".","shuffle",":","Flag","indicating","whether","apply","shuffle","when","dividing","elements","between","workers","shuffle_seed",":","Random","seed","used","to","shuffle","when","dividing","elements","across","workers","generator",":","torch.Generator","(",")","object","used","as","a","source","of","randomness","when","selecting","items","for","the","next","round","on","a","given","worker"]
code: def __init__(
      self,
      *,
      total_size: int,
      sample_rate: float,
      shuffle: bool = True,
      shuffle_seed: int = 0,
      steps: int = None,
      generator=None,
  ):
      """
  
      Args:
          total_size: total number of samples to sample from
          sample_rate: number of samples to draw.
          shuffle: Flag indicating whether apply shuffle when dividing elements
              between workers
          shuffle_seed: Random seed used to shuffle when dividing elements across workers
          generator: torch.Generator() object used as a source of randomness
              when selecting items for the next round on a given worker
      """
      self.total_size = total_size
      self.sample_rate = sample_rate
      self.generator = generator
      self.num_replicas = torch.distributed.get_world_size()
      self.rank = torch.distributed.get_rank()
      self.epoch = 0
      self.shuffle = shuffle
      self.shuffle_seed = shuffle_seed
  
      if self.total_size <= 0:
          raise ValueError(
              "total_size should be a positive integer "
              "value, but got total_size={}".format(self.total_size)
          )
  
      # Size of the local dataset specific to the current replica
      self.num_samples = self.total_size // self.num_replicas
      if self.rank < self.total_size % self.num_replicas:
          # The first replicas get an extra datapoint if necessary (balanced)
          self.num_samples += 1
  
      # Number of batches: same as non-distributed Poisson sampling, but each batch is smaller
      if steps is not None:
          self.num_batches = steps
      else:
          self.num_batches = int(1 / self.sample_rate)
  
file_path: opacus/opacus/utils/uniform_sampler.py
imports: from typing import List
  import torch
  from torch.utils.data import Sampler
label_desc: Use image node_id 1 to create a new DistributedUniformWithReplacementSampler object from inherited base classes: Sampler with example: obj = DistributedUniformWithReplacementSampler()