class_docstr: DataLoader subclass that always does Poisson sampling and supports empty batches
  by default.
  
  Typically instantiated via ``DPDataLoader.from_data_loader()`` method based
  on another DataLoader. DPDataLoader would preserve the behaviour of the original
  data loader, except for the two aspects.
  
  First, it switches ``batch_sampler`` to ``UniformWithReplacementSampler``, thus enabling
  Poisson sampling (i.e. each element in the dataset is selected to be in the
  next batch with a certain probability defined by ``sample_rate`` parameter).
  NB: this typically leads to a batches of variable size.
  NB2: By default, ``sample_rate`` is calculated based on the ``batch_size`` of the
  original data loader, so that the average batch size stays the same
  
  Second, it wraps collate function with support for empty batches.
  Most PyTorch modules will happily process tensors of shape ``(0, N, ...)``,
  but many collate functions will fail to produce such a batch. As with the
  Poisson sampling empty batches become a possibility, we need a DataLoader that
  can handle them.
class_docstr_tok: ["DataLoader","subclass","that","always","does","Poisson","sampling","and","supports","empty","batches","by","default",".","Typically","instantiated","via","``","DPDataLoader.from_data_loader","(",")","``","method","based","on","another","DataLoader",".","DPDataLoader","would","preserve","the","behaviour","of","the","original","data","loader",",","except","for","the","two","aspects",".","First",",","it","switches","``","batch_sampler","``","to","``","UniformWithReplacementSampler","``",",","thus","enabling","Poisson","sampling","(","i.e",".","each","element","in","the","dataset","is","selected","to","be","in","the","next","batch","with","a","certain","probability","defined","by","``","sample_rate","``","parameter",")",".","NB",":","this","typically","leads","to","a","batches","of","variable","size",".","NB2",":","By","default",",","``","sample_rate","``","is","calculated","based","on","the","``","batch_size","``","of","the","original","data","loader",",","so","that","the","average","batch","size","stays","the","same","Second",",","it","wraps","collate","function","with","support","for","empty","batches",".","Most","PyTorch","modules","will","happily","process","tensors","of","shape","``","(","0",",","N",",","...",")","``",",","but","many","collate","functions","will","fail","to","produce","such","a","batch",".","As","with","the","Poisson","sampling","empty","batches","become","a","possibility",",","we","need","a","DataLoader","that","can","handle","them","."]
docstr: Args:
      dataset: See :class:`torch.utils.data.DataLoader`
      sample_rate: probability with which each element of the dataset is included
          in the next batch.
      num_workers: See :class:`torch.utils.data.DataLoader`
      collate_fn: See :class:`torch.utils.data.DataLoader`
      pin_memory: See :class:`torch.utils.data.DataLoader`
      drop_last: See :class:`torch.utils.data.DataLoader`
      timeout: See :class:`torch.utils.data.DataLoader`
      worker_init_fn: See :class:`torch.utils.data.DataLoader`
      multiprocessing_context: See :class:`torch.utils.data.DataLoader`
      generator: Random number generator used to sample elements
      prefetch_factor: See :class:`torch.utils.data.DataLoader`
      persistent_workers: See :class:`torch.utils.data.DataLoader`
      distributed: set ``True`` if you'll be using DPDataLoader in a DDP environment
          Selects between ``DistributedUniformWithReplacementSampler`` and
          ``UniformWithReplacementSampler`` sampler implementations
docstr_tok: ["Args",":","dataset",":","See",":","class",":","`","torch.utils.data.DataLoader","`","sample_rate",":","probability","with","which","each","element","of","the","dataset","is","included","in","the","next","batch",".","num_workers",":","See",":","class",":","`","torch.utils.data.DataLoader","`","collate_fn",":","See",":","class",":","`","torch.utils.data.DataLoader","`","pin_memory",":","See",":","class",":","`","torch.utils.data.DataLoader","`","drop_last",":","See",":","class",":","`","torch.utils.data.DataLoader","`","timeout",":","See",":","class",":","`","torch.utils.data.DataLoader","`","worker_init_fn",":","See",":","class",":","`","torch.utils.data.DataLoader","`","multiprocessing_context",":","See",":","class",":","`","torch.utils.data.DataLoader","`","generator",":","Random","number","generator","used","to","sample","elements","prefetch_factor",":","See",":","class",":","`","torch.utils.data.DataLoader","`","persistent_workers",":","See",":","class",":","`","torch.utils.data.DataLoader","`","distributed",":","set","``","True","``","if","you","'ll","be","using","DPDataLoader","in","a","DDP","environment","Selects","between","``","DistributedUniformWithReplacementSampler","``","and","``","UniformWithReplacementSampler","``","sampler","implementations"]
code: def __init__(
      self,
      dataset: Dataset,
      *,
      sample_rate: float,
      collate_fn: Optional[_collate_fn_t] = None,
      drop_last: bool = False,
      generator=None,
      distributed: bool = False,
      **kwargs,
  ):
      """
  
      Args:
          dataset: See :class:`torch.utils.data.DataLoader`
          sample_rate: probability with which each element of the dataset is included
              in the next batch.
          num_workers: See :class:`torch.utils.data.DataLoader`
          collate_fn: See :class:`torch.utils.data.DataLoader`
          pin_memory: See :class:`torch.utils.data.DataLoader`
          drop_last: See :class:`torch.utils.data.DataLoader`
          timeout: See :class:`torch.utils.data.DataLoader`
          worker_init_fn: See :class:`torch.utils.data.DataLoader`
          multiprocessing_context: See :class:`torch.utils.data.DataLoader`
          generator: Random number generator used to sample elements
          prefetch_factor: See :class:`torch.utils.data.DataLoader`
          persistent_workers: See :class:`torch.utils.data.DataLoader`
          distributed: set ``True`` if you'll be using DPDataLoader in a DDP environment
              Selects between ``DistributedUniformWithReplacementSampler`` and
              ``UniformWithReplacementSampler`` sampler implementations
      """
  
      self.sample_rate = sample_rate
      self.distributed = distributed
  
      if distributed:
          batch_sampler = DistributedUniformWithReplacementSampler(
              total_size=len(dataset),  # type: ignore[assignment, arg-type]
              sample_rate=sample_rate,
              generator=generator,
          )
      else:
          batch_sampler = UniformWithReplacementSampler(
              num_samples=len(dataset),  # type: ignore[assignment, arg-type]
              sample_rate=sample_rate,
              generator=generator,
          )
      sample_empty_shapes = [(0, *shape_safe(x)) for x in dataset[0]]
      dtypes = [dtype_safe(x) for x in dataset[0]]
      if collate_fn is None:
          collate_fn = default_collate
  
      if drop_last:
          logger.warning(
              "Ignoring drop_last as it is not compatible with DPDataLoader."
          )
  
      super().__init__(
          dataset=dataset,
          batch_sampler=batch_sampler,
          collate_fn=wrap_collate_with_empty(
              collate_fn=collate_fn,
              sample_empty_shapes=sample_empty_shapes,
              dtypes=dtypes,
          ),
          generator=generator,
          **kwargs,
      )
  
file_path: opacus/opacus/data_loader.py
imports: import logging
  from functools import partial
  from typing import Any, List, Optional, Sequence, Tuple, Type, Union
  import torch
  from opacus.utils.uniform_sampler import DistributedUniformWithReplacementSampler, UniformWithReplacementSampler
  from torch.utils.data import BatchSampler, DataLoader, Dataset, IterableDataset, Sampler
  from torch.utils.data._utils.collate import default_collate
  from torch.utils.data.dataloader import _collate_fn_t
label_desc: Use image node_id 1 to create a new DPDataLoader object from inherited base classes: DataLoader with example: obj = DPDataLoader(dataset)