How to use WeightedRandomSampler in PyTorch

Imbalanced training data can make a model update mostly from the majority class even when every row is valid. PyTorch WeightedRandomSampler changes which map-style dataset indices a DataLoader draws, allowing underrepresented samples to appear more often during training.

The sampler expects one non-negative weight per dataset row rather than one weight per class. Inverse class frequencies become per-sample weights only after the class weight tensor is indexed by the label assigned to each row.

Apply the sampler only to the training loader so validation and test metrics retain their original class distribution. Sampling with replacement permits repeated minority rows, while a dedicated torch.Generator makes a smoke test repeatable without changing the project's global random state.

1. Define a map-style training dataset and its label tensor in weighted_sampler_loader.py.

   weighted_sampler_loader.py

   import torch
   from torch.utils.data import DataLoader, TensorDataset, WeightedRandomSampler
    
    
   labels = torch.tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1])
   features = torch.arange(len(labels), dtype=torch.float32).unsqueeze(1)
   train_dataset = TensorDataset(features, labels)

   Project data can replace features and labels when the label order matches train_dataset[i]. WeightedRandomSampler applies only to map-style datasets with indexable rows.

2. Append the inverse-frequency sample weights below the dataset definition.

   class_counts = torch.bincount(labels)
   class_weights = 1.0 / class_counts.float()
   sample_weights = class_weights[labels]

   sample_weights has one entry per training row. The weights express relative draw probabilities and do not need to sum to 1.

3. Configure the weighted sampler and training loader below the weight calculation.

   sampler = WeightedRandomSampler(
       weights=sample_weights,
       num_samples=len(sample_weights),
       replacement=True,
       generator=torch.Generator().manual_seed(7),
   )
   train_loader = DataLoader(train_dataset, batch_size=4, sampler=sampler)

   shuffle=True conflicts with an explicit sampler because DataLoader treats those arguments as mutually exclusive.

4. Add a fail-capable class-exposure check below the loader construction.

   sampled_labels = torch.cat(
       [batch_labels for _, batch_labels in train_loader]
   )
   sampled_counts = torch.bincount(sampled_labels, minlength=len(class_counts))
    
   print(f"original_counts={class_counts.tolist()}")
   print(f"sampled_counts={sampled_counts.tolist()}")
   print(f"minority_share={sampled_counts[1].item() / len(sampled_labels):.2f}")
    
   assert len(sampled_labels) == len(labels)
   assert sampled_counts[1] > class_counts[1]

   The fixed seed makes this smoke-test result repeatable. Random training epochs can differ from an exact class balance even when their expected sampling probabilities are equal.

5. Run the completed weighted-loader smoke test.

   $ python weighted_sampler_loader.py
   original_counts=[9, 3]
   sampled_counts=[5, 7]
   minority_share=0.58

   The sampled epoch keeps 12 draws while increasing class 1 from 3 source rows to 7 sampled appearances. A wrong weight order, sampler attachment, or replacement setting can make either assertion fail.