title: Using WeightedRandomSampler in PyTorch
body: <p>I need to implement a multi-label image classification model in PyTorch. However my data is not balanced, so I used the <code>WeightedRandomSampler</code> in PyTorch to create a custom dataloader. But when I iterate through the custom dataloader, I get the error : <code>IndexError: list index out of range</code> </p>
  
  <p>Implemented the following code using this link :<a href="https://discuss.pytorch.org/t/balanced-sampling-between-classes-with-torchvision-dataloader/2703/3?u=surajsubramanian" rel="noreferrer">https://discuss.pytorch.org/t/balanced-sampling-between-classes-with-torchvision-dataloader/2703/3?u=surajsubramanian</a></p>
  
  <pre class="lang-py prettyprint-override"><code>def make_weights_for_balanced_classes(images, nclasses):                        
      count = [0] * nclasses                                                      
      for item in images:                                                         
          count[item[1]] += 1                                                     
      weight_per_class = [0.] * nclasses                                      
      N = float(sum(count))                                                   
      for i in range(nclasses):                                                   
          weight_per_class[i] = N/float(count[i])                                 
      weight = [0] * len(images)                                              
      for idx, val in enumerate(images):                                          
          weight[idx] = weight_per_class[val[1]]                                  
      return weight 
  </code></pre>
  
  <pre class="lang-py prettyprint-override"><code>weights = make_weights_for_balanced_classes(train_dataset.imgs, len(full_dataset.classes))
  weights = torch.DoubleTensor(weights)
  sampler = WeightedRandomSampler(weights, len(weights))
  
  train_loader = DataLoader(train_dataset, batch_size=4,sampler = sampler, pin_memory=True)   
  </code></pre>
  
  <p>Based on the answer in <a href="https://stackoverflow.com/a/60813495/10077354">https://stackoverflow.com/a/60813495/10077354</a>, the following is my updated code. But then too when I create a dataloader :<code>loader = DataLoader(full_dataset, batch_size=4, sampler=sampler)</code>, <code>len(loader)</code> returns 1.</p>
  
  <pre class="lang-py prettyprint-override"><code>class_counts = [1691, 743, 2278, 1271]
  num_samples = np.sum(class_counts)
  labels = [tag for _,tag in full_dataset.imgs] 
  
  class_weights = [num_samples/class_counts[i] for i in range(len(class_counts)]
  weights = [class_weights[labels[i]] for i in range(num_samples)]
  sampler = WeightedRandomSampler(torch.DoubleTensor(weights), num_samples)
  </code></pre>
  
  <p>Thanks a lot in advance !</p>
  
  <p>I included an utility function based on the accepted answer below :</p>
  
  <pre class="lang-py prettyprint-override"><code>def sampler_(dataset):
      dataset_counts = imageCount(dataset)
      num_samples = sum(dataset_counts)
      labels = [tag for _,tag in dataset]
  
      class_weights = [num_samples/dataset_counts[i] for i in range(n_classes)]
      weights = [class_weights[labels[i]] for i in range(num_samples)]
      sampler = WeightedRandomSampler(torch.DoubleTensor(weights), int(num_samples))
      return sampler
  </code></pre>
  
  <p>The imageCount function finds number of images of each class in the dataset. Each row in the dataset contains the image and the class, so we take the second element in the tuple into consideration.</p>
  
  <pre class="lang-py prettyprint-override"><code>def imageCount(dataset):
      image_count = [0]*(n_classes)
      for img in dataset:
          image_count[img[1]] += 1
      return image_count
  </code></pre>