🎓 A1Learner
🔍

PyTorch Class Imbalance

Class imbalance means one class holds far more rows than the other. A model can score high by favouring the big class and ignoring the small one.

The trap in counts

Two hundred rows hold only ten sick cases. A model that finds nobody still scores 0.95 accuracy.

y = torch.tensor([0] * 190 + [1] * 10)
pred = torch.zeros(200, dtype=torch.long)
acc = (pred == y).float().mean().item()
print(f"Accuracy: {acc:.2f}")
print("Sick found:", pred[y == 1].sum().item())

Set the positive weight

pos_weight scales the loss on positive rows. Nineteen matches the ratio of 190 healthy rows to 10 sick rows.

loss_fn = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([19.0]))
print(loss_fn.pos_weight)

Example

import torch
from torch import nn

def run(use_weights):
    torch.manual_seed(0)
    n0, n1 = 190, 10
    x0 = torch.randn(n0, 2) + torch.tensor([-0.5, -0.5])
    x1 = torch.randn(n1, 2) + torch.tensor([0.5, 0.5])
    x = torch.cat([x0, x1])
    y = torch.cat([torch.zeros(n0), torch.ones(n1)]).reshape(-1, 1)
    order = torch.randperm(len(x))
    x, y = x[order], y[order]
    torch.manual_seed(1)
    model = nn.Linear(2, 1)
    if use_weights:
        loss_fn = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([19.0]))
    else:
        loss_fn = nn.BCEWithLogitsLoss()
    opt = torch.optim.Adam(model.parameters(), lr=0.05)
    for epoch in range(40):
        loss = loss_fn(model(x), y)
        opt.zero_grad()
        loss.backward()
        opt.step()
    model.eval()
    with torch.no_grad():
        pred = (model(x) > 0).float()
    sick = y == 1
    recall = (pred[sick] == 1).float().mean().item()
    acc = (pred == y).float().mean().item()
    return acc, recall

a0, r0 = run(False)
a1, r1 = run(True)
print(f"No weights:   accuracy {a0:.2f}, sick recall {r0:.2f}")
print(f"With weights: accuracy {a1:.2f}, sick recall {r1:.2f}")

Key points

  • Minority class: the class with far fewer rows
  • A model that always picks the majority class still scores the majority share as accuracy
  • Class weight: a multiplier that makes minority mistakes cost more in the loss
  • Use pos_weight in BCEWithLogitsLoss to weight the positive class
  • Recall on the minority class shows whether the weighting worked

In this example, one function trains the same problem twice. The data holds 190 healthy rows and 10 sick rows in two overlapping clouds. A shuffle mixes the classes, and the seed resets so both runs start alike. The first run uses a plain BCEWithLogitsLoss. It reaches accuracy 0.94 and finds only 0.20 of the sick rows. The second run sets pos_weight to 19, the class ratio. Its accuracy drops to 0.84, and its sick recall climbs to 0.90. The two lines print the trade directly.

Note

Weighting trades some overall accuracy for minority recall, and the closing example shows both numbers

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

What does pos_weight do in BCEWithLogitsLoss?