class_docstr: ClassPowerlawSampler
  
  Batch-sampler that samples 'batch-size' images using a power law distribution
  Intended to be used on balanced dataset such as Imagenet
  
  :param world_size: total number of workers in network
  :param rank: local rank in network
  :param batch_size: number of images to load from each class
  :param powerlaw: strength of the power law (coefficient)
  :param seed: common seed across workers for subsampling classes
class_docstr_tok: ["ClassPowerlawSampler","Batch-sampler","that","samples","'batch-size","'","images","using","a","power","law","distribution","Intended","to","be","used","on","balanced","dataset","such","as","Imagenet",":","param","world_size",":","total","number","of","workers","in","network",":","param","rank",":","local","rank","in","network",":","param","batch_size",":","number","of","images","to","load","from","each","class",":","param","powerlaw",":","strength","of","the","power","law","(","coefficient",")",":","param","seed",":","common","seed","across","workers","for","subsampling","classes"]
code: def __init__(
      self,
      dataset,
      num_classes: int,
      class_index: Dict[int, List[int]],
      labels: List[int],
      world_size: int,
      rank: int,
      batch_size=1,
      powerlaw=1.0,
      seed=0,
  ):
      super(ClassPowerlawSampler, self).__init__(dataset)
  
      self.dataset = dataset
      self.num_classes = num_classes
      self.class_index = class_index
      self.world_size = world_size
      self.rank = rank
      self.batch_size = batch_size
      self.class_indices = []
      self.set_epoch(0)
  
      g = torch.Generator()
      g.manual_seed(seed)
      if powerlaw < 0.0:
          # HACK to plug inverse square root re-sampling - replace by proper sampler
          probs_map = self.inverse_square_root_resampling(labels)
          probs = np.array(
              [probs_map[c] for c in range(self.num_classes)]
          )
          self.probs = probs
      else:
          probs = np.array(
              [
                  1.0 / (1.0 + i) ** powerlaw
                  for i in range(self.num_classes)
              ]
          )
          probs /= sum(probs)
          self.probs = probs[torch.randperm(len(probs), generator=g)]
  
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 ClassPowerlawSampler object from inherited base classes: torch.utils.data with example: obj = ClassPowerlawSampler(dataset, num_classes, class_index, labels, world_size, rank, batch_size, powerlaw, seed)