How to use WeightedRandomSampler in PyTorch

Class imbalance can make a PyTorch training loop see the majority class far more often than the minority class. WeightedRandomSampler changes the index order used by a DataLoader so selected samples appear with probabilities based on weights instead of their raw dataset frequency.

The sampler receives one weight for each sample, not one weight for each class. A common pattern is to count labels, invert those class counts, and then index the class weights with the training labels to build the per-sample weight vector.

Use the weighted sampler on the training loader for a map-style dataset, where each integer index maps to one sample. Keep validation and test loaders unweighted so their metrics still reflect the original data distribution.

Steps to use PyTorch WeightedRandomSampler:

Create a weighted sampler smoke script.

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)
dataset = TensorDataset(features, labels)
 
class_counts = torch.bincount(labels)
class_weights = 1.0 / class_counts.float()
sample_weights = class_weights[labels]
 
generator = torch.Generator().manual_seed(42)
sampler = WeightedRandomSampler(
    weights=sample_weights,
    num_samples=24,
    replacement=True,
    generator=generator,
)
loader = DataLoader(dataset, batch_size=6, sampler=sampler)
 
sampled_labels = []
batch_labels = []
for _, label_batch in loader:
    batch = label_batch.tolist()
    batch_labels.append(batch)
    sampled_labels.extend(batch)
 
sampled_counts = torch.bincount(torch.tensor(sampled_labels), minlength=2)
 
print(f"torch_version={torch.__version__}")
print(f"original_counts={class_counts.tolist()}")
print(f"original_minority_share={class_counts[1].item() / len(labels):.2f}")
print(f"class_weights={[round(value, 4) for value in class_weights.tolist()]}")
print(f"sample_weights_first_last={[round(sample_weights[0].item(), 4), round(sample_weights[-1].item(), 4)]}")
print(f"sampled_batches={batch_labels}")
print(f"sampled_counts={sampled_counts.tolist()}")
print(f"sampled_minority_share={sampled_counts[1].item() / len(sampled_labels):.2f}")
print(f"minority_original={class_counts[1].item()}")
print(f"minority_sampled={sampled_counts[1].item()}")
print(f"minority_oversampled={sampled_counts[1].item() > class_counts[1].item()}")

The smoke script draws 24 samples from 12 rows so the changed class exposure is easy to see. For a normal training epoch, set num_samples to len(sample_weights) unless a longer sampled epoch is intentional.

Run the smoke script.

$ python weighted_sampler_loader.py
torch_version=2.12.1+cpu
original_counts=[9, 3]
original_minority_share=0.25
class_weights=[0.1111, 0.3333]
sample_weights_first_last=[0.1111, 0.3333]
sampled_batches=[[0, 0, 0, 0, 1, 0], [1, 1, 0, 0, 0, 0], [1, 1, 0, 1, 0, 1], [1, 0, 0, 1, 0, 0]]
sampled_counts=[15, 9]
sampled_minority_share=0.38
minority_original=3
minority_sampled=9
minority_oversampled=True

minority_oversampled=True confirms that class 1 appeared more often in the sampled loader than it did in the original label list.

Collect one label for each training sample in dataset order.
```
train_labels = torch.as_tensor(train_dataset.targets, dtype=torch.long)
```
Replace train_dataset.targets with the label list or metadata field used by the dataset. The order must match the indices returned by train_dataset[i].
Convert class counts into per-sample weights.
```
class_counts = torch.bincount(train_labels)
class_weights = 1.0 / class_counts.float()
sample_weights = class_weights[train_labels]
```
sample_weights must have the same length as the training dataset. WeightedRandomSampler treats these as sample probabilities after internal normalization, so the weights do not need to sum to 1.
Create the weighted sampler for the training epoch.
```
sampler = WeightedRandomSampler(
    weights=sample_weights,
    num_samples=len(sample_weights),
    replacement=True,
)
```
replacement=True allows minority-class samples to be drawn more than once in an epoch. Pass a torch.Generator when repeated runs need the same sample order.
Related: How to set a random seed in PyTorch
Attach the sampler to the training DataLoader.
```
train_loader = DataLoader(
    train_dataset,
    batch_size=64,
    sampler=sampler,
    num_workers=4,
)
```
Do not also set shuffle=True on this loader. The sampler is already controlling the index order for the map-style training dataset.
Check the sampled training labels before a full training run.
```
sampled_labels = []
for _, batch_labels in train_loader:
    sampled_labels.extend(batch_labels.tolist())
 
print(torch.bincount(torch.tensor(sampled_labels), minlength=len(class_counts)))
```
Adjust the batch unpacking when the dataset returns dictionaries or more than two fields. The printed counts should show more minority-class exposure than the raw training labels.

Use the weighted loader in the training loop.

for features, labels in train_loader:
    optimizer.zero_grad(set_to_none=True)
    predictions = model(features)
    loss = loss_fn(predictions, labels)
    loss.backward()
    optimizer.step()

Remove the smoke script after the loader pattern is copied into the project.
```
$ rm weighted_sampler_loader.py
```