Debugging Training: Curves, Overfitting & Sanity Checks

"My loss would not budge for six hours. They blamed the learning rate, the optimizer, my very architecture. The labels were shifted by one. I had been perfectly, diligently learning to predict the wrong answer the entire time."

A Network That Was Never Actually Broken

You now have the full recipe from the previous five sections. This section is what to do when, despite all of it, the run misbehaves, which it will. The central reframing is diagnostic: training problems are not mysterious, they leave signatures in the curves and respond predictably to cheap probes. The most important habit is to assume a bug before assuming the model is fundamentally incapable, because the prior probability of a pipeline bug (a wrong label mapping, an un-shuffled loader, a forgotten normalization from Section 21.1, a frozen layer that should not be) vastly exceeds the probability that a well-known architecture simply cannot learn your task.

1. Reading the Curves: Three Canonical Shapes Beginner

The plot of training loss and validation loss against epochs is the single most informative diagnostic you have, and it has three readable shapes. Underfitting: both losses are high and roughly equal, and both are still falling or have plateaued at a poor value, meaning the model lacks the capacity, training time, or learning rate to fit even the training data. Overfitting: training loss keeps falling toward zero while validation loss bottoms out and then rises, the classic divergence that means the model is memorizing the training set at the expense of generalization. A healthy run: both losses fall together and the validation loss flattens near, but slightly above, the training loss. Figure 21.6.1 contrasts the three.

Each shape prescribes a different fix. For underfitting, increase capacity, train longer, raise the learning rate, or reduce regularization and augmentation, because the model needs to fit harder. For overfitting, do the opposite: add the augmentation of Section 21.2, increase weight decay or dropout (the random unit-masking regularizer from Chapter 19), get more data, or simply early-stop at the validation minimum. A healthy run needs no intervention except to let it finish. The mistake to avoid is reaching for a random knob; the curve already told you which direction to move. The same loss-and-accuracy curves return for the detection and segmentation tasks of Chapter 23 and Chapter 24, where the validation metric is mean average precision (mAP) or mean intersection over union (mIoU) rather than top-1 accuracy but the three shapes read identically.

2. The Single Most Valuable Sanity Check: Overfit One Batch Beginner

Before launching a full run, do this every time: take a single small batch and train on it repeatedly until the model fits it perfectly, driving the loss to near zero. A model with enough capacity should memorize a handful of examples easily within a few dozen steps. If it cannot, your training is broken in a way that no amount of full-dataset training will fix, and you have found out in thirty seconds instead of six hours. This one check catches an enormous fraction of real bugs: a detached gradient, a learning rate of zero, a label-shape mismatch, a loss that is not connected to the parameters, a frozen layer that should be trainable.

import torch

def overfit_one_batch(model, batch, optimizer, criterion, steps=200):
    """A model that cannot drive ONE batch to ~0 loss has a bug, not a data problem."""
    images, labels = batch
    model.train()
    for step in range(steps):
        optimizer.zero_grad()
        loss = criterion(model(images), labels)
        loss.backward()
        optimizer.step()
        if step % 50 == 0:
            print(f"step {step:3d}  loss {loss.item():.4f}")
    return loss.item()

# Expected on a healthy setup (e.g. ResNet-18, batch of 16):
# step   0  loss 2.3107
# step  50  loss 0.0431
# step 100  loss 0.0009
# step 150  loss 0.0002      <- near zero: the optimization path works
# If the loss is stuck near 2.30 (= ln 10 for 10 classes), the gradient is not flowing.

3. Look at the Data the Model Actually Sees Intermediate

A startling fraction of training bugs are invisible in code and obvious in a picture. The discipline is to pull one batch straight from the data loader, after all transforms, and actually look at it: render the images, print the labels, and confirm they match. This single habit surfaces the augmentation-gone-wrong from Section 21.2 (a crop so aggressive the object is gone), the normalization error from Section 21.1 (images that look like noise because the channel order is swapped, the same RGB-versus-BGR channel-order trap first flagged in Chapter 0), and the off-by-one label shift from this section's epigraph (the model was learning perfectly, just to predict the neighbor's label), the exact failure the illustration below dramatizes. The code below dumps a labeled grid.

import torch, torchvision

def inspect_batch(loader, class_names, mean, std):
    """Render one real, post-transform batch with its labels. Look at it."""
    images, labels = next(iter(loader))
    # Undo normalization so the saved image is human-viewable.
    m = torch.tensor(mean).view(1, 3, 1, 1)
    s = torch.tensor(std).view(1, 3, 1, 1)
    vis = (images * s + m).clamp(0, 1)
    grid = torchvision.utils.make_grid(vis[:16], nrow=4)
    torchvision.utils.save_image(grid, "batch_check.png")
    print("labels:", [class_names[i] for i in labels[:16].tolist()])
    print("pixel range:", vis.min().item(), "to", vis.max().item())   # should be ~0..1
    print("batch shape:", tuple(images.shape))

# Open batch_check.png and confirm: do the images look right, and do the
# printed labels match what you SEE in each tile? Most data bugs die here.

4. A Playbook for "It Will Not Learn" Intermediate

When a model refuses to learn, work the checks in order of cost, cheapest first, so you spend probe time in proportion to how likely and how cheap each cause is. The ordered playbook below is the distilled habit. It deliberately puts the data and wiring checks before any hyperparameter tuning, because a bug there makes all downstream tuning meaningless, exactly the trap the practical example fell into.

Two recurring culprits deserve a direct mention because they are common and easily missed. An exploding or NaN loss almost always means the learning rate is too high or you forgot gradient clipping; halve the rate or clip the gradient norm (rescale the gradient vector so its overall magnitude never exceeds a set ceiling, which caps any single runaway update). A loss that is flat from step zero usually means the gradient is not reaching the parameters at all, a detached tensor, a layer left in evaluation mode, or requires_grad mistakenly left False from the freezing of Section 21.3. Both show up immediately in the overfit-one-batch check, which is why it sits so high in the playbook. The Hands-On Lab at the end of this section is where the whole chapter comes together: you will assemble the data, augmentation, transfer, recipe, and sanity-check links into one script that trains and evaluates a real classifier on a small custom set.

import lightning as L

trainer = L.Trainer(
    max_epochs=100,
    overfit_batches=1,        # built-in: overfit ONE batch to validate the pipeline
    fast_dev_run=False,       # set True to run 1 train+val batch end-to-end as a smoke test
    log_every_n_steps=10,     # automatic loss/metric logging for the curves of subsection 1
)
# trainer.fit(model, train_loader, val_loader)
# Curves stream to TensorBoard; overfit_batches=1 reproduces subsection 2 with one flag.

pip install torch torchvision timm matplotlib

import torch
from torchvision import transforms
from torchvision.datasets import OxfordIIITPet
from torch.utils.data import DataLoader

MEAN, STD = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]

train_tf = transforms.Compose([
    transforms.RandomResizedCrop(224, scale=(0.6, 1.0)),
    transforms.RandomHorizontalFlip(),
    transforms.RandAugment(num_ops=2, magnitude=9),   # automated policy, Section 21.2
    transforms.ToTensor(),
    transforms.Normalize(MEAN, STD),
])
# TODO: build val_tf with a deterministic Resize(256) + CenterCrop(224),
#       then ToTensor and the SAME Normalize(MEAN, STD). No randomness here.
val_tf = ...

train = OxfordIIITPet("./data", split="trainval", download=True, transform=train_tf)
val   = OxfordIIITPet("./data", split="test",     download=True, transform=val_tf)
train_loader = DataLoader(train, batch_size=32, shuffle=True,  num_workers=0)
val_loader   = DataLoader(val,   batch_size=64, shuffle=False, num_workers=0)
print(f"train={len(train)}  val={len(val)}  classes={len(train.classes)}")

import timm

device = "cuda" if torch.cuda.is_available() else "cpu"
# TODO: create a pretrained "resnet50" with num_classes set to the dataset's
#       class count, then move it to `device`. Hint: timm.create_model(...).
model = ...
print("head out features:", model.get_classifier().out_features)  # expect 37

import math
import torch.nn as nn

xb, yb = next(iter(train_loader))
xb, yb = xb.to(device), yb.to(device)

# Loss at init should be about ln(K): the chance level for K classes.
with torch.no_grad():
    init_loss = nn.CrossEntropyLoss()(model(xb), yb).item()
print(f"loss at init = {init_loss:.3f}  (expected about {math.log(len(train.classes)):.3f})")

# Overfit ONE batch: a healthy model drives it toward zero in well under 100 steps.
probe = nn.CrossEntropyLoss()
opt_probe = torch.optim.Adam(model.parameters(), lr=1e-3)
model.train()
for step in range(60):
    opt_probe.zero_grad()
    loss = probe(model(xb), yb)
    loss.backward(); opt_probe.step()
    # TODO: every 20 steps, print step and loss.item() so you can watch it fall.
    ...

from torch.optim.lr_scheduler import LinearLR, CosineAnnealingLR, SequentialLR

model = timm.create_model("resnet50", pretrained=True,
                          num_classes=len(train.classes)).to(device)  # fresh after probe
EPOCHS = 8
steps_per_epoch = len(train_loader)
warmup_steps = steps_per_epoch  # one warmup epoch

opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.05)
# TODO: build a SequentialLR that does LinearLR warmup for `warmup_steps`,
#       then CosineAnnealingLR for the remaining steps. Milestone = [warmup_steps].
sched = ...
loss_fn = nn.CrossEntropyLoss(label_smoothing=0.1)   # label smoothing, Section 21.4

import numpy as np

def mixup(x, y, num_classes, alpha=0.2):
    lam = np.random.beta(alpha, alpha)
    idx = torch.randperm(x.size(0), device=x.device)
    mixed_x = lam * x + (1 - lam) * x[idx]
    y1 = nn.functional.one_hot(y, num_classes).float()
    mixed_y = lam * y1 + (1 - lam) * y1[idx]   # mix the LABELS too
    return mixed_x, mixed_y

for epoch in range(EPOCHS):
    model.train()
    for xb, yb in train_loader:
        xb, yb = xb.to(device), yb.to(device)
        mx, my = mixup(xb, yb, len(train.classes))
        opt.zero_grad()
        # Soft-target cross-entropy: -sum(target * log_softmax(logits)).
        logp = nn.functional.log_softmax(model(mx), dim=1)
        loss = -(my * logp).sum(dim=1).mean()
        loss.backward()
        opt.step()
        sched.step()   # TODO: confirm this is per-STEP, not per-epoch
    print(f"epoch {epoch}  last train loss {loss.item():.3f}  lr {sched.get_last_lr()[0]:.2e}")

@torch.no_grad()
def accuracy(model, loader):
    model.eval()
    correct = total = 0
    for xb, yb in loader:
        xb, yb = xb.to(device), yb.to(device)
        preds = model(xb).argmax(dim=1)
        # TODO: accumulate correct predictions and the total count
        ...
    return 100.0 * correct / total

print(f"validation top-1 accuracy: {accuracy(model, val_loader):.2f}%")

import matplotlib.pyplot as plt
# Assume you appended to `train_losses` and `val_accs` (one value per epoch) above.
fig, ax1 = plt.subplots(figsize=(6, 4))
ax1.plot(train_losses, color="tab:blue", label="train loss")
ax1.set_xlabel("epoch"); ax1.set_ylabel("train loss", color="tab:blue")
ax2 = ax1.twinx()
ax2.plot(val_accs, color="tab:orange", label="val acc")
ax2.set_ylabel("val accuracy (%)", color="tab:orange")
plt.title("Recipe training curve"); plt.tight_layout()
plt.savefig("recipe_curve.png", dpi=120)
print("saved recipe_curve.png")

import math
import numpy as np
import torch
import torch.nn as nn
import timm
import matplotlib.pyplot as plt
from torch.optim.lr_scheduler import LinearLR, CosineAnnealingLR, SequentialLR
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.datasets import OxfordIIITPet

device = "cuda" if torch.cuda.is_available() else "cpu"
MEAN, STD = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]

# Step 1: data with augmentation + normalization
train_tf = transforms.Compose([
    transforms.RandomResizedCrop(224, scale=(0.6, 1.0)),
    transforms.RandomHorizontalFlip(),
    transforms.RandAugment(num_ops=2, magnitude=9),
    transforms.ToTensor(),
    transforms.Normalize(MEAN, STD),
])
val_tf = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(MEAN, STD),
])
train = OxfordIIITPet("./data", split="trainval", download=True, transform=train_tf)
val   = OxfordIIITPet("./data", split="test",     download=True, transform=val_tf)
K = len(train.classes)
train_loader = DataLoader(train, batch_size=32, shuffle=True,  num_workers=0)
val_loader   = DataLoader(val,   batch_size=64, shuffle=False, num_workers=0)

# Step 3: sanity checks (on a throwaway model)
probe_model = timm.create_model("resnet50", pretrained=True, num_classes=K).to(device)
xb, yb = next(iter(train_loader))
xb, yb = xb.to(device), yb.to(device)
with torch.no_grad():
    print(f"loss at init = {nn.CrossEntropyLoss()(probe_model(xb), yb).item():.3f} "
          f"(expected about {math.log(K):.3f})")
opt_probe = torch.optim.Adam(probe_model.parameters(), lr=1e-3)
probe_model.train()
for step in range(60):
    opt_probe.zero_grad()
    loss = nn.CrossEntropyLoss()(probe_model(xb), yb)
    loss.backward(); opt_probe.step()
    if step % 20 == 0:
        print(f"overfit step {step:3d} loss {loss.item():.3f}")
del probe_model, opt_probe

# Step 2 + 4: fresh model and the modern recipe
model = timm.create_model("resnet50", pretrained=True, num_classes=K).to(device)
EPOCHS = 8
steps_per_epoch = len(train_loader)
warmup_steps = steps_per_epoch
opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.05)
sched = SequentialLR(
    opt,
    schedulers=[
        LinearLR(opt, start_factor=0.01, total_iters=warmup_steps),
        CosineAnnealingLR(opt, T_max=EPOCHS * steps_per_epoch - warmup_steps),
    ],
    milestones=[warmup_steps],
)

# Step 5: MixUp training loop
def mixup(x, y, num_classes, alpha=0.2):
    lam = np.random.beta(alpha, alpha)
    idx = torch.randperm(x.size(0), device=x.device)
    mixed_x = lam * x + (1 - lam) * x[idx]
    y1 = nn.functional.one_hot(y, num_classes).float()
    mixed_y = lam * y1 + (1 - lam) * y1[idx]
    return mixed_x, mixed_y

@torch.no_grad()
def accuracy(model, loader):
    model.eval()
    correct = total = 0
    for xb, yb in loader:
        xb, yb = xb.to(device), yb.to(device)
        preds = model(xb).argmax(dim=1)
        correct += (preds == yb).sum().item()
        total += yb.size(0)
    return 100.0 * correct / total

train_losses, val_accs = [], []
for epoch in range(EPOCHS):
    model.train()
    for xb, yb in train_loader:
        xb, yb = xb.to(device), yb.to(device)
        mx, my = mixup(xb, yb, K)
        opt.zero_grad()
        logp = nn.functional.log_softmax(model(mx), dim=1)
        loss = -(my * logp).sum(dim=1).mean()
        loss.backward()
        opt.step()
        sched.step()
    acc = accuracy(model, val_loader)
    train_losses.append(loss.item()); val_accs.append(acc)
    print(f"epoch {epoch}  train loss {loss.item():.3f}  "
          f"val acc {acc:.2f}%  lr {sched.get_last_lr()[0]:.2e}")

# Step 6 + 7: final accuracy and curve
print(f"final validation top-1 accuracy: {val_accs[-1]:.2f}%")
fig, ax1 = plt.subplots(figsize=(6, 4))
ax1.plot(train_losses, color="tab:blue", label="train loss")
ax1.set_xlabel("epoch"); ax1.set_ylabel("train loss", color="tab:blue")
ax2 = ax1.twinx()
ax2.plot(val_accs, color="tab:orange", label="val acc")
ax2.set_ylabel("val accuracy (%)", color="tab:orange")
plt.title("Recipe training curve"); plt.tight_layout()
plt.savefig("recipe_curve.png", dpi=120)
print("saved recipe_curve.png")