Introduction

WeightedRandomSampler in PyTorch lets you sample dataset items according to custom probabilities. It is commonly used for imbalanced datasets, but the crucial detail is that the sampler expects one weight per sample, not one weight per class.

That distinction is the source of most first-time mistakes. If your dataset has ten thousand examples, the sampler needs ten thousand weights.

The Mental Model

For class imbalance, the usual process is:

1. count how many samples belong to each class
2. compute a class weight such as the inverse frequency
3. map each class weight onto every sample belonging to that class
4. pass the full sample-level weight vector to WeightedRandomSampler

So if class 0 is common and class 1 is rare, samples from class 1 receive larger weights and are drawn more often.

A Runnable Example

Here is a small imbalanced dataset with 90 examples from class 0 and 10 from class 1:

1import torch
2from torch.utils.data import TensorDataset, DataLoader, WeightedRandomSampler
3
4labels = torch.tensor([0] * 90 + [1] * 10)
5features = torch.randn(len(labels), 4)
6
7dataset = TensorDataset(features, labels)
8
9class_counts = torch.bincount(labels)
10class_weights = 1.0 / class_counts.float()
11sample_weights = class_weights[labels]
12
13sampler = WeightedRandomSampler(
14    weights=sample_weights,
15    num_samples=len(sample_weights),
16    replacement=True
17)
18
19loader = DataLoader(dataset, batch_size=16, sampler=sampler)
20
21batch_features, batch_labels = next(iter(loader))
22print(batch_labels)

The important line is this one:

sample_weights = class_weights[labels]

That converts a short class-weight vector into the per-sample weight vector the sampler actually needs.

Why replacement=True Is Common

For oversampling minority classes, replacement=True is usually the intended setting. It allows the same rare sample to be drawn more than once within an epoch, which is how oversampling actually happens.

If you set replacement=False, each sample can appear at most once in a pass. That can still be useful, but it reduces the balancing effect because minority examples cannot be repeated.

A simple rule is:

• use replacement=True when you want minority classes to appear more often
• use replacement=False only when repeated sampling would be undesirable

Understanding num_samples

num_samples controls how many draws the sampler makes in one pass through the DataLoader. A common choice is the dataset length:

num_samples = len(sample_weights)

That produces an epoch with as many draws as there are original samples, but the class distribution of those draws is influenced by the weights.

You can choose a different value if you want longer or shorter effective epochs, but keeping it equal to the dataset size is the least surprising default.

Sampler Versus Loss Weighting

Balanced sampling and weighted loss are different tools. WeightedRandomSampler changes which examples the model sees and how often it sees them. A class-weighted loss changes how strongly each training example affects the optimization step.

You can combine them, but doing both means you are compensating for imbalance twice. That may be useful in some situations, but it should be a conscious decision rather than an automatic habit.

Common Pitfalls

The biggest mistake is passing class weights directly into WeightedRandomSampler. If your dataset has many samples, the sampler needs a weight vector of that same length.

Another pitfall is expecting every mini-batch to be perfectly balanced. The sampler improves balance over repeated draws, but small individual batches can still fluctuate.

A third issue is evaluating the model on oversampled data. Resampling is a training trick, not a substitute for realistic validation. Keep validation and test loaders representative of the real distribution.

Summary

• 'WeightedRandomSampler expects one weight per dataset sample.'
• For class imbalance, compute class weights and then expand them into sample weights.
• 'replacement=True is common when you want minority classes to be oversampled.'
• 'num_samples controls the effective epoch length.'
• Use balanced sampling for training, but evaluate on a realistic validation set.