The Training Loop: Losses, Metrics & Checkpointing

"Every epoch I do the same five things in the same order. People call it a loop. I call it a discipline, and it is the only reason any of you trust the number at the end."

A Training Loop That Never Skips a Step

We now have all the parts. Section 18.1 gave us a model, Section 18.2 the optimizer and loss, Section 18.3 the PyTorch abstractions, and Section 18.4 the data stream. The training loop is the conductor that brings them together. It is short, perhaps forty lines, and it is the single most important code in this entire book to understand cold, because the convolutional networks of Chapter 19, the transformers of Chapter 22, and even the diffusion models of Chapter 33 are all trained by this same loop with a different model and loss swapped in.

1. The Five-Step Core Beginner

At the heart of training is a five-step ritual performed once per batch, in this exact order. It is the same five-beat loop Section 18.3 named, zero, forward, loss, backward, step, now written out in full. Zero the gradients (because they accumulate, as Section 18.3 warned). Run the forward pass to get predictions. Compute the scalar loss. Run loss.backward() to fill every parameter's gradient. Call optimizer.step() to update the parameters. Figure 18.5.1 lays out the full loop, with this five-beat core nested inside the epoch and phase structure.

2. The Complete Loop in Code Intermediate

The code below is the whole loop, assembled from the chapter's parts and runnable end to end. It separates the train and validation phases, toggles model.train() and model.eval() as Section 18.3's practical example taught, wraps validation in torch.no_grad(), tracks both loss and accuracy, and saves the model whenever validation accuracy reaches a new best. Read it as the reference implementation: every training script in Part III is a variation on it.

# The complete canonical training loop: run_epoch executes both the train and
# validation phases via a single flag, and train_model wires the optimizer and
# schedule, checkpoints the best validation accuracy, and stops early on a stall.
import torch
import torch.nn as nn

def run_epoch(model, loader, criterion, optimizer, device, train):
    model.train() if train else model.eval()
    total_loss, correct, n = 0.0, 0, 0
    with torch.set_grad_enabled(train):      # build the graph only in the train phase; restored on exit
        for x, y in loader:
            x, y = x.to(device), y.to(device)    # move batch to GPU
            if train:
                optimizer.zero_grad()            # 1. reset accumulated gradients
            logits = model(x)                    # 2. forward pass
            loss = criterion(logits, y)          # 3. scalar loss
            if train:
                loss.backward()                  # 4. fill .grad on every parameter
                optimizer.step()                 # 5. update parameters
            total_loss += loss.item() * x.size(0)
            correct += (logits.argmax(1) == y).sum().item()
            n += x.size(0)
    return total_loss / n, correct / n       # average loss, accuracy

def train_model(model, train_loader, val_loader, epochs=20,
                lr=3e-4, device="cuda", patience=5):
    model.to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
    best_acc, stale = 0.0, 0
    for epoch in range(epochs):
        tr_loss, tr_acc = run_epoch(model, train_loader, criterion, optimizer, device, train=True)
        va_loss, va_acc = run_epoch(model, val_loader,   criterion, optimizer, device, train=False)
        scheduler.step()                     # advance the LR schedule
        if va_acc > best_acc:                # checkpoint the best model only
            best_acc, stale = va_acc, 0
            torch.save({"model": model.state_dict(), "epoch": epoch,
                        "val_acc": va_acc}, "best.pt")
        else:
            stale += 1
            if stale >= patience:            # early stopping
                print(f"early stop at epoch {epoch}; best val acc {best_acc:.4f}")
                break
        print(f"epoch {epoch:02d}  train {tr_acc:.4f}  val {va_acc:.4f}")
    return best_acc

A few details earn their place. loss.item() pulls a Python float off the GPU tensor and detaches it from the graph (calling it inside the loop without .item() would silently retain the entire graph and leak memory). Multiplying by x.size(0) before summing correctly weights the final partial batch when averaging. And the checkpoint saves the state_dict (the parameter tensors) rather than the whole model object, which is the portable, version-robust way PyTorch recommends; you reload it into a freshly constructed model with model.load_state_dict(torch.load("best.pt")["model"]).

3. Choosing the Loss and the Metric Beginner

The loss is what the optimizer minimizes; the metric is what you actually care about, and they are often different. For multi-class classification the loss is cross-entropy (Section 18.2); for multi-label problems it is binary cross-entropy per label; for regression it is mean squared or L1 error. The metric, by contrast, is chosen for human interpretability and business relevance: accuracy when classes are balanced, but per-class precision, recall, and F1 when they are not, exactly the lesson from the imbalanced-data practical example in Section 18.4. You minimize cross-entropy because it is smooth and differentiable, but you report F1 because that is what tells you whether the defect detector works. As the book proceeds, the metrics specialize, mean average precision for object detection in Chapter 23 and intersection-over-union for segmentation in Chapter 24, but the loss-versus-metric distinction is constant.

4. Checkpointing, Resuming, and Early Stopping Intermediate

Two failure modes motivate checkpointing. First, the best model is rarely the last one: a run that trains for 20 epochs may peak at epoch 12 and degrade afterward, so saving only the final weights throws away the best result, which is why the loop above saves on every validation improvement. Second, long runs crash, machines preempt, and spot instances vanish; a periodic full checkpoint (model, optimizer state, scheduler state, and epoch number) lets you resume exactly where you left off rather than restarting from scratch. The optimizer state matters because Adam's momentum buffers (Section 18.2) are part of the training state; resuming without them restarts the optimizer cold. Early stopping, halting when validation has not improved for a patience window, both saves compute and acts as a regularizer by preventing the overfitting phase, and the loop above implements it in four lines.