docstr: ClassStratifiedSampler
  Batch-sampler that samples 'batch-size' images from subset of randomly
  chosen classes e.g., if classes a,b,c are randomly sampled,
  the sampler returns
      torch.cat([a,b,c], [a,b,c], ..., [a,b,c], dim=0)
  where a,b,c, are images from classes a,b,c respectively.
  Sampler, samples images WITH REPLACEMENT (i.e., not epoch-based)
  
  :param world_size: total number of workers in network
  :param rank: local rank in network
  :param batch_size: num. images to load from each class
  :param classes_per_batch: num. classes to randomly sample for batch
  :param epochs: num consecutive epochs thru data_source before gen.reset
  :param seed: common seed across workers for subsampling classes
  :param unique_classes: true ==> each worker samples a distinct set of classes; false ==> all workers sample the same classes
docstr_tok: ["ClassStratifiedSampler","Batch-sampler","that","samples","'batch-size","'","images","from","subset","of","randomly","chosen","classes","e.g.",",","if","classes","a",",","b",",","c","are","randomly","sampled",",","the","sampler","returns","torch.cat","(","[","a",",","b",",","c","]",",","[","a",",","b",",","c","]",",","...",",","[","a",",","b",",","c","]",",","dim=0",")","where","a",",","b",",","c",",","are","images","from","classes","a",",","b",",","c","respectively",".","Sampler",",","samples","images","WITH","REPLACEMENT","(","i.e.",",","not","epoch-based",")",":","param","world_size",":","total","number","of","workers","in","network",":","param","rank",":","local","rank","in","network",":","param","batch_size",":","num",".","images","to","load","from","each","class",":","param","classes_per_batch",":","num",".","classes","to","randomly","sample","for","batch",":","param","epochs",":","num","consecutive","epochs","thru","data_source","before","gen.reset",":","param","seed",":","common","seed","across","workers","for","subsampling","classes",":","param","unique_classes",":","true","==",">","each","worker","samples","a","distinct","set","of","classes",";","false","==",">","all","workers","sample","the","same","classes"]
code: def __init__(
      self,
      dataset,
      num_classes: int,
      labels: np.array,
      class_index: Dict[int, List[int]],
      world_size: int,
      rank: int,
      batch_size=1,
      classes_per_batch=10,
      epochs=1,
      seed=0,
      unique_classes=False,
  ):
      """
      ClassStratifiedSampler
      Batch-sampler that samples 'batch-size' images from subset of randomly
      chosen classes e.g., if classes a,b,c are randomly sampled,
      the sampler returns
          torch.cat([a,b,c], [a,b,c], ..., [a,b,c], dim=0)
      where a,b,c, are images from classes a,b,c respectively.
      Sampler, samples images WITH REPLACEMENT (i.e., not epoch-based)
  
      :param world_size: total number of workers in network
      :param rank: local rank in network
      :param batch_size: num. images to load from each class
      :param classes_per_batch: num. classes to randomly sample for batch
      :param epochs: num consecutive epochs thru data_source before gen.reset
      :param seed: common seed across workers for subsampling classes
      :param unique_classes: true ==> each worker samples a distinct set of classes; false ==> all workers sample the same classes
      """
      super().__init__(dataset)
  
      self.rank = rank
      self.world_size = world_size
      self.cpb = classes_per_batch
      self.unique_cpb = unique_classes
      self.batch_size = batch_size
      self.num_classes = num_classes
      self.labels = labels  # Used for debugging only
      self.class_index = class_index
      self.epochs = epochs
      self.epoch = 0
  
      if not self.unique_cpb:
          assert self.num_classes % self.cpb == 0
  
      self.base_seed = seed  # instance seed
      self.seed = seed
  
imports: import contextlib
  import logging
  import queue
  from typing import Dict, List
  import numpy
  import torch
  from PIL import Image
  from torch.utils.data import Dataset
  from torch.utils.data.distributed import DistributedSampler
label_desc: Use image node_id 1 to create a new StratifiedClassSampler object from inherited base classes: torch.utils.data with example: obj = StratifiedClassSampler(dataset, num_classes, labels, class_index, world_size, rank, batch_size, classes_per_batch, epochs, seed, unique_classes)