class_docstr: BatchSampler - from a MNIST-like dataset, samples batch_size according to given input distribution
  assuming multi-class labels
  adapted from https://github.com/adambielski/siamese-triplet/blob/master/datasets.py
class_docstr_tok: ["BatchSampler","-","from","a","MNIST-like","dataset",",","samples","batch_size","according","to","given","input","distribution","assuming","multi-class","labels","adapted","from","https",":","\/\/github.com\/adambielski\/siamese-triplet\/blob\/master\/datasets.py"]
code: def __init__(self, dataset, batch_size, class_weights):
      labels = get_labels(dataset)
      self._classes = sorted(set(labels))
  
      n_classes = len(self._classes)
      if n_classes > len(class_weights):
          k = len(class_weights)
          sum_w = np.sum(class_weights)
          if sum_w >= 1:
              # normalize attributing equal weight to weighted part and remaining part
              class_weights /= (
                  sum_w * k / n_classes + (n_classes - k) / n_classes
              )
          krem = k - n_classes
          wrem = 1 - sum_w
          logging.warning(
              f"will assume uniform distribution for labels > {len(class_weights)}"
          )
          self._class_weights = np.ones(n_classes, dtype=np.float64)
          self._class_weights[:k] = class_weights
          self._class_weights[k:] = wrem / krem
      else:
          self._class_weights = class_weights[:n_classes]
  
      if np.sum(self._class_weights) != 1:
          self._class_weights = self._class_weights / np.sum(
              self._class_weights
          )
  
      logging.debug("Using weights=", self._class_weights)
      if batch_size == 0:
          raise ValueError(
              f"batch_size should be bigger than the number of classes, got {batch_size}"
          )
  
      self._class_to_iter = {
          class_: InfiniteSliceIterator(
              np.where(labels == class_)[0], class_=class_
          )
          for class_ in self._classes
      }
  
      self.n_dataset = len(labels)
      self._batch_size = batch_size
      self._n_batches = self.n_dataset // self._batch_size
      if self._n_batches == 0:
          raise ValueError(
              f"Dataset is not big enough to generate batches with size {self._batch_size}"
          )
      logging.debug("K=", n_classes, "nk=", self._batch_size)
      logging.debug("Batch size = ", self._batch_size)
  
imports: import logging
  import numpy
  import torch.utils.data
  import torchvision
  from torch.utils.data.sampler import BatchSampler, RandomSampler
label_desc: Use image node_id 1 to create a new ReweightedBatchSampler object from inherited base classes: torch.utils.data.sampler with example: obj = ReweightedBatchSampler(dataset, batch_size, class_weights)