🎓 A1Learner
🔍

PyTorch Focal Loss

Focal loss is a loss that shrinks the score of examples the model already gets right. Hard examples keep their full score, so they count more in training.

The shrinking factor

An easy example has pt near 1. A hard example has pt far below 1.

import torch

p_easy = torch.tensor(0.95)
p_hard = torch.tensor(0.30)
print(f"easy factor {(1 - p_easy) ** 2:.4f}")
print(f"hard factor {(1 - p_hard) ** 2:.4f}")

Two examples under both losses

The first logit is confident and correct. The second logit is wrong. BCE is binary cross-entropy, the per-example loss.

import torch
import torch.nn.functional as F

logits = torch.tensor([2.9444, -0.8473])
targets = torch.tensor([1.0, 1.0])
bce = F.binary_cross_entropy_with_logits(logits, targets, reduction="none")
p = torch.sigmoid(logits)
focal = (1 - p) ** 2 * bce
print(f"bce: {bce[0].item():.4f} {bce[1].item():.4f}")
print(f"focal: {focal[0].item():.4f} {focal[1].item():.4f}")

Gamma sets the strength

The same easy example shrinks more as gamma grows. At gamma 0 nothing shrinks.

import torch

pt = torch.tensor(0.9)
for gamma in [0, 1, 2]:
    print(f"gamma {gamma}: factor {(1 - pt) ** gamma:.2f}")

Example

import torch

torch.manual_seed(0)

# Class 0 is common and tight. Class 1 is rare and spread out.
n0, n1 = 380, 20
X0 = torch.randn(n0, 2) * 0.7 + torch.tensor([-2.0, -2.0])
X1 = torch.randn(n1, 2) * 1.6 + torch.tensor([0.5, 0.5])
X = torch.cat([X0, X1])
y = torch.cat([torch.zeros(n0), torch.ones(n1)]).reshape(-1, 1)


def focal_loss(logits, targets, gamma=2.0):
    bce = torch.nn.functional.binary_cross_entropy_with_logits(
        logits, targets, reduction="none"
    )
    p = torch.sigmoid(logits)
    pt = torch.where(targets == 1, p, 1 - p)
    return (((1 - pt) ** gamma) * bce).mean()


def train(use_focal):
    torch.manual_seed(1)
    model = torch.nn.Linear(2, 1)
    optimizer = torch.optim.Adam(model.parameters(), lr=0.05)
    for step in range(100):
        logits = model(X)
        if use_focal:
            loss = focal_loss(logits, y)
        else:
            loss = torch.nn.functional.binary_cross_entropy_with_logits(logits, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    with torch.no_grad():
        pred = (torch.sigmoid(model(X)) >= 0.5).float()
        hard = y == 1
        recall = (pred[hard] == 1).float().mean().item()
    return recall


print(f"Examples: {len(X)} (class 0: {n0}, class 1: {n1})")
print(f"Cross-entropy hard-class recall: {train(False):.2f}")
print(f"Focal hard-class recall: {train(True):.2f}")

Key points

  • Imbalanced data: one class appears far more often than another
  • Easy example: an example the model already predicts with high confidence
  • pt: the model's probability for the correct class on one example
  • Use the factor (1 - pt) ** gamma to shrink the loss of easy examples
  • Gamma: the number that sets how hard easy examples are shrunk, often 2
  • Recall: the share of a class's examples that the model finds
  • Focal loss leaves hard examples at full strength and turns easy examples down

In this example, the data holds 400 points. Class 0 has 380 points in a tight cluster. Class 1 has 20 points spread wide, so its points are the hard ones. Two models start from the same seed and train for 100 steps. The cross-entropy model finds 45 percent of the hard class, a recall of 0.45. The focal model finds 55 percent, a recall of 0.55. The easy class-0 points stop dominating the focal loss, so the rare class gets more of the training signal.

Note

With gamma set to 0, focal loss is exactly ordinary cross-entropy

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

An example has pt = 0.99 under gamma 2. What happens to its loss under focal loss?