Training Loops, Checkpoints, and Early Stopping

A PyTorch training loop is the control system around a model. It decides when tensors move through the network, when gradients are accumulated, when parameters change, when validation is measured, and which model state is worth keeping. Checkpoints and early stopping extend that loop so training can recover from interruption and stop when additional epochs are no longer improving generalization.

In this part of the Deep Learning with PyTorch course, the model architecture is only one actor. The loop is where the experiment becomes reproducible: it connects DataLoader batches, nn.Module modes, loss functions, optimizers, schedulers, metrics, files, and stopping policy. The outcome is a loop you can resume, inspect, and defend when validation behavior changes.

How the Loop Works

A basic supervised loop repeats the same sequence for each mini-batch. The model runs a forward pass and produces predictions. A loss function compares predictions with targets and returns a scalar tensor. loss.backward() asks autograd to traverse the recorded computation graph and write gradients into each leaf parameter’s .grad buffer. The optimizer reads those buffers and updates parameter tensors in-place. Finally, the next batch starts after gradients are cleared with optimizer.zero_grad().

Two details matter internally. First, PyTorch accumulates gradients by default. If you forget to clear them, the second batch’s gradient is added to the first batch’s gradient, which is useful for deliberate gradient accumulation but wrong for ordinary mini-batch SGD. Second, modules have training and evaluation modes. model.train() enables training behavior such as dropout randomness and batch normalization statistic updates. model.eval() disables those training-specific behaviors, but it does not turn off gradient tracking. Validation should usually combine model.eval() with torch.no_grad() so results reflect inference behavior and avoid building unused graphs.

A checkpoint is a serialized snapshot of enough state to continue or select training. For normal PyTorch modules, the portable core is model.state_dict(), a mapping from parameter and buffer names to tensors. Resume checkpoints also include optimizer.state_dict(), the current epoch or step, the best validation metric so far, early-stopping counters, and often scheduler, scaler, random-number-generator, or configuration state. Saving only the model weights is enough for inference; it is not enough to resume AdamW with the same momentum estimates.

Early stopping watches a validation metric across epochs. It defines what counts as improvement, how much non-improvement is tolerated, and which saved state will be restored. For loss, lower is better; for accuracy or F1, higher is better. patience is the number of consecutive validation checks allowed without meaningful improvement. min_delta prevents tiny noise from resetting patience. The best practice is to checkpoint the best validation state as soon as it appears, not the final state at the moment stopping triggers.

API Anatomy

The common loop vocabulary is compact. for epoch in range(...) controls full passes over training data. for xb, yb in train_loader controls mini-batches. optimizer.zero_grad(set_to_none=True) clears gradients efficiently by setting buffers to None. loss.backward() computes gradients. optimizer.step() mutates parameters. torch.save(obj, path) writes checkpoint objects, and torch.load(path, map_location=...) reads them. load_state_dict copies tensors back into a module or optimizer.

The ordering is part of the contract. A typical batch is zero gradients, forward, loss, backward, optional gradient clipping, optimizer step, and optional scheduler step. Validation happens after training batches for an epoch, inside eval and no_grad. Checkpointing happens after validation because the decision should be based on held-out behavior, not the training loss from the most recent mini-batch.

Example 1: One Batch Update

This first example isolates the minimum train step. The deterministic output proves the loss is a Python float when extracted with item(), and that gradients were written for the layer’s weight tensor before the optimizer stepped.

import torch
from torch import nn

# A single-batch training step: forward pass, scalar loss, backward pass, optimizer update.
torch.manual_seed(7)
model = nn.Linear(3, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
loss_fn = nn.MSELoss()

features = torch.tensor([[1.0, 0.0, 2.0], [0.0, 1.0, -1.0]])
targets = torch.tensor([[1.0], [-1.0]])

optimizer.zero_grad(set_to_none=True)
predictions = model(features)
loss = loss_fn(predictions, targets)
loss.backward()
optimizer.step()

print(type(loss.item()).__name__)
print(model.weight.grad.shape)

Expected output is float followed by torch.Size([1, 3]). If the gradient shape does not match the parameter shape, the model definition or loss path is not the one you think it is. Notice that the code does not call model.eval(); this is a training batch, so training mode is the right default.

Example 2: Save Resume State

A checkpoint is usually a dictionary because training state has several owners. This example writes and reads a checkpoint in a temporary file. The important behavior is not the filename; it is the presence of model weights, optimizer state, the epoch cursor, best validation loss, and the early-stopping counter.

import tempfile
import torch
from torch import nn

torch.manual_seed(3)
model = nn.Linear(2, 1)
optimizer = torch.optim.AdamW(model.parameters(), lr=0.01)
checkpoint = {
    "epoch": 4,
    "model_state": model.state_dict(),
    "optimizer_state": optimizer.state_dict(),
    "best_val_loss": 0.318,
    "patience_used": 1,
}

with tempfile.NamedTemporaryFile(suffix=".pt") as tmp:
    torch.save(checkpoint, tmp.name)
    restored = torch.load(tmp.name, weights_only=False)

print(sorted(restored.keys()))
print(restored["epoch"], restored["best_val_loss"])

The expected first line lists ['best_val_loss', 'epoch', 'model_state', 'optimizer_state', 'patience_used']. The second line is 4 0.318. In a real run, loading would call model.load_state_dict(restored["model_state"]) and optimizer.load_state_dict(restored["optimizer_state"]) before continuing at the next epoch.

Example 3: Patience and Minimum Delta

Early stopping is a small state machine. It stores the best metric and the count of consecutive checks that failed to improve enough. In this example, a loss drop smaller than 0.01 is treated as noise.

class EarlyStopper:
    def __init__(self, patience: int, min_delta: float = 0.0):
        self.patience = patience
        self.min_delta = min_delta
        self.best = float("inf")
        self.bad_epochs = 0

    def step(self, val_loss: float) -> bool:
        if val_loss < self.best - self.min_delta:
            self.best = val_loss
            self.bad_epochs = 0
            return False
        self.bad_epochs += 1
        return self.bad_epochs >= self.patience

stopper = EarlyStopper(patience=2, min_delta=0.01)
for epoch, val_loss in enumerate([0.90, 0.82, 0.811, 0.810, 0.809], start=1):
    if stopper.step(val_loss):
        print(f"stop at epoch {epoch}; best={stopper.best:.3f}")
        break

The output is stop at epoch 4; best=0.820. Epoch 3 has loss 0.811, which is not strictly less than 0.82 - 0.01, so it does not reset patience. This strict comparison is a design choice; changing it to <= changes stopping behavior on ties.

Design Choices and Trade-Offs

Validation frequency trades signal quality against cost. Validating every epoch is clear for small datasets. For very large datasets, validating every fixed number of steps or on a representative validation subset may be cheaper, but it can make early stopping noisier. Batch size trades throughput against gradient noise and memory pressure. Gradient accumulation can simulate a larger batch, but then the loop must step the optimizer only after the chosen number of mini-batches.

Checkpoint strategy also has trade-offs. Keeping only the latest checkpoint limits disk use but loses the best validation model if the latest epoch overfits. Keeping every epoch improves auditability but can waste storage quickly. A common compromise is to keep last.pt for resume and best.pt for selection. For long jobs, write checkpoints atomically by saving to a temporary path and renaming it when the write succeeds, so an interruption does not leave a half-written file with the expected name.

Metric direction must be explicit. A loop written for validation loss will make the wrong decision if someone swaps in accuracy without changing the comparison. Store the metric name, direction, and dataset split in the checkpoint metadata. That metadata makes the file understandable months later and prevents accidentally comparing training loss from one run with validation loss from another.

Failure Modes and Troubleshooting

Symptom: training loss behaves erratically or explodes after the first few batches. Cause: gradients may be accumulating unintentionally. Diagnose: print or log gradient norms before optimizer.step() and inspect whether zero_grad is called once per intended optimizer update. Correct: call optimizer.zero_grad(set_to_none=True) at the start of each update, or implement deliberate accumulation with a documented accumulation counter.

Symptom: validation loss is much better during training than after loading the model for inference. Cause: validation may have run in training mode, leaving dropout active or batch normalization updating. Diagnose: assert not model.training inside validation and run the same validation batch twice; dropout should not cause different outputs in eval mode. Correct: call model.eval() before validation and return to model.train() before the next training epoch.

Symptom: resumed training diverges from the original run after a crash. Cause: the checkpoint contains weights but not optimizer, scheduler, scaler, or epoch state. Adam-family optimizers depend on internal moment estimates, so resuming with fresh optimizer state changes the trajectory. Diagnose: inspect checkpoint keys and compare optimizer state dict sizes before and after loading. Correct: save all state required for continuation, not just inference weights.

Symptom: early stopping halts too soon on a noisy validation curve. Cause: patience is shorter than the metric noise pattern, or the validation set is too small. Diagnose: plot validation metric by epoch and mark best checkpoints. Correct: increase patience, use a larger validation set, smooth only for reporting, or choose a metric with lower variance.

Reliability and Performance Implications

Reliable loops make state transitions explicit. Use deterministic seeds for debugging, but do not promise perfect reproducibility across all devices and kernels unless you have configured deterministic algorithms and accepted the possible performance cost. On GPU, move model and batch tensors to the same device before the forward pass. For validation and inference, torch.no_grad() reduces memory use because autograd does not retain intermediate tensors.

Checkpoint files can contain more than tensors when saved with Python pickle-based serialization, so load checkpoints only from trusted sources. For model distribution, prefer documented weight files and load them into known code. Store checkpoints outside source control when they are large or contain sensitive derived information from private datasets.

Hands-On Lab: Build a Stoppable Loop

Prerequisites: Python with PyTorch installed, a writable scratch directory, and enough CPU memory for a tiny synthetic regression dataset. No external data is required.

1. Create a new script named train_with_stopping.py in a scratch directory.
2. Paste the following complete loop. It creates data, trains a small network, evaluates on held-out samples, tracks the best validation loss, and stops after three non-improving epochs.
3. Run the script once and confirm both printed values are True.
4. Add a checkpoint write when validation improves: save best, model_state, optimizer_state, and epoch.
5. Restart from that file by loading the model and optimizer states, then continue from epoch + 1.

import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

torch.manual_seed(11)
x = torch.linspace(-1, 1, 80).unsqueeze(1)
y = 2 * x + 0.1 * torch.sin(8 * x)
train_ds = TensorDataset(x[:60], y[:60])
val_ds = TensorDataset(x[60:], y[60:])
train_loader = DataLoader(train_ds, batch_size=15, shuffle=True)
val_loader = DataLoader(val_ds, batch_size=20)

model = nn.Sequential(nn.Linear(1, 8), nn.Tanh(), nn.Linear(8, 1))
optimizer = torch.optim.AdamW(model.parameters(), lr=0.05)
loss_fn = nn.MSELoss()
best = {"loss": float("inf"), "state": None, "epoch": 0}
patience = 3
bad_epochs = 0

for epoch in range(1, 31):
    model.train()
    for xb, yb in train_loader:
        optimizer.zero_grad(set_to_none=True)
        loss = loss_fn(model(xb), yb)
        loss.backward()
        optimizer.step()

    model.eval()
    total = 0.0
    count = 0
    with torch.no_grad():
        for xb, yb in val_loader:
            batch_loss = loss_fn(model(xb), yb)
            total += batch_loss.item() * xb.size(0)
            count += xb.size(0)
    val_loss = total / count

    if val_loss < best["loss"] - 1e-4:
        best = {"loss": val_loss, "state": model.state_dict(), "epoch": epoch}
        bad_epochs = 0
    else:
        bad_epochs += 1
    if bad_epochs >= patience:
        break

print(best["epoch"] > 0)
print(best["loss"] < 1.0)

Verification: the first printed value confirms that at least one epoch became the best epoch. The second confirms the best validation loss crossed a loose quality threshold for this simple regression task. Also verify that your checkpoint contains the epoch cursor and optimizer state, not only model weights.

Cleanup: remove the scratch script and generated checkpoint files when finished. If you added checkpoints to a project repository by mistake, remove the large binary files and add the checkpoint directory to the project’s ignore rules before committing.

Assessment Exercises

1. Modify the early-stopping class so it can maximize accuracy instead of minimizing loss. What state and comparison logic must change?
2. A loop validates every 10,000 steps rather than every epoch. Explain how this changes the meaning of patience and checkpoint names.
3. You load best.pt and get missing keys from load_state_dict. List two plausible causes and the diagnostic command or inspection you would use first.
4. Design a checkpoint retention policy for a three-day GPU training job with limited disk. Which files do you keep, and why?
5. Where should torch.no_grad() appear in the loop, and what incorrect behavior would result from placing it around the training forward pass?

Summary

Training loops are the procedural core of PyTorch learning. They coordinate modes, gradients, optimizer updates, validation measurements, saved state, and stopping decisions. A strong loop saves both the latest resume state and the best validation state, evaluates in inference mode, records enough metadata to understand the decision, and stops based on a metric policy that matches the task. Those details turn a model definition into a repeatable experiment.