🎓 A1Learner
🔍

PyTorch Weighted Sampling

Weighted sampling draws rows with unequal chances, so rare rows appear more often. A WeightedRandomSampler takes one weight per row and draws by those weights.

The imbalance itself

torch.bincount counts the rows per label. Class 0 holds ninety rows and class 2 holds two, so plain sampling mostly serves class 0.

labels = torch.tensor([0] * 90 + [1] * 8 + [2] * 2)
counts = torch.bincount(labels, minlength=3)
print("Rows per class:", counts.tolist())

One weight per row

The weight is one over the class count. Rare rows weigh more, and each class then weighs the same in total.

class_weights = 1.0 / counts.float()
print("Class weights:", [round(float(w), 4) for w in class_weights])
row_weights = class_weights[labels]
print("Weight of a class 0 row:", round(float(row_weights[0]), 4))
print("Weight of a class 2 row:", round(float(row_weights[-1]), 4))

A seeded generator repeats the draws

Most drawn indices point at the rare rows at the end of the dataset. The same seed repeats the same draws.

def first_draws():
    gen = torch.Generator().manual_seed(7)
    sampler = WeightedRandomSampler(
        weights=row_weights.double(), num_samples=10, replacement=True, generator=gen
    )
    return list(sampler)

print("First run: ", first_draws())
print("Second run:", first_draws())
print("Identical:", first_draws() == first_draws())

Example

import torch
from torch.utils.data import WeightedRandomSampler

torch.manual_seed(0)

# 100 rows. Class 0 is common. Classes 1 and 2 are rare.
labels = torch.tensor([0] * 90 + [1] * 8 + [2] * 2)

counts = torch.bincount(labels, minlength=3)
print("Rows per class:", counts.tolist())

draws = 300

plain_gen = torch.Generator().manual_seed(7)
plain = WeightedRandomSampler(
    weights=torch.ones(len(labels), dtype=torch.double),
    num_samples=draws,
    replacement=True,
    generator=plain_gen,
)
plain_counts = torch.bincount(labels[list(plain)], minlength=3)
print("Drawn without weights:", plain_counts.tolist())

# One weight per row. Rare rows weigh more, so each class weighs the same in total.
class_weights = 1.0 / counts.float()
row_weights = class_weights[labels]
weighted_gen = torch.Generator().manual_seed(7)
weighted = WeightedRandomSampler(
    weights=row_weights.double(),
    num_samples=draws,
    replacement=True,
    generator=weighted_gen,
)
weighted_counts = torch.bincount(labels[list(weighted)], minlength=3)
print("Drawn with weights:", weighted_counts.tolist())

Key points

  • WeightedRandomSampler: a sampler that draws each row with a chance set by its weight
  • Use one weight per row, often one over the row's class count
  • Replacement: a drawn row goes back, so it can be drawn again
  • Pass a seeded generator so the draws repeat exactly

In this example, the labels describe 100 rows. Class 0 holds 90 rows, class 1 holds 8, and class 2 holds 2. The first sampler gives every row weight 1, so 300 draws land close to the row counts. Class 0 takes 263 draws, and class 2 takes 7. The second sampler weights each row by one over its class count. The same 300 draws now spread across the classes. Class 0 takes 99, class 1 takes 90, and class 2 takes 111. Both samplers use a seeded generator, so the counts repeat exactly on every run.

Note

Weighted sampling changes which rows training sees, but it does not create new data.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

A dataset holds 90 rows of class 0 and 2 rows of class 2. What weight should class 2 rows get, relative to class 0 rows?