04 · Best Model & Early Stopping 🏅

4 min read

The model at the last epoch is rarely the best model. This sub-module is how you capture the best one — and stop wasting time once it stops improving.

1. The 'Why'

Module 05 ended with an uncomfortable truth revealed by the training curves: as training continues, the model's performance on the validation set improves for a while, then plateaus, then often gets worse as the model begins to overfit. This means the model you happen to have when the loop finishes — the last-epoch model — is frequently not the best one you trained. Somewhere in the middle was a version that generalized better, and if you only saved at the end, you threw it away. Best-model checkpointing fixes this: you watch a validation metric every epoch and save a checkpoint only when it improves, so at the end you still hold the single best version regardless of how the run finished.

Its natural companion is early stopping. If validation performance hasn't improved for several epochs, continuing is usually a waste of compute and risks drifting further into overfitting. Early stopping watches for that stagnation and halts training automatically once a "patience" window passes with no improvement. Together these two techniques turn the raw training loop from Module 05 into a disciplined procedure that reliably produces — and preserves — a good model. They lean on everything from this module (state_dict saving, checkpoints) and close the arc that began with reading train/validation curves.

2. Core Concepts

Best-model checkpointing. Track the best validation score seen so far. Each epoch, if the current score beats it, save the model's state_dict (overwriting best_model.pt) and update the record. At the end, that file holds the best model.

Which metric to monitor. Pick one validation metric as the criterion — lower is better for loss, higher is better for accuracy. Be consistent about the direction.

Early stopping & "patience". Count epochs since the last improvement. When that counter reaches your patience threshold, break out of the training loop. Patience trades a little extra training against the risk of stopping too soon on a noisy plateau.

Restore the best at the end. After the loop (whether it finished or stopped early), load best_model.pt back into the model so you're using the best version, not the last.

3. Code in Action

Saving the best model each epoch

import torch

best_val = float("inf")            # we're monitoring validation LOSS (lower is better)

for epoch in range(100):
    train_loss = train_one_epoch(model, train_loader, criterion, optimizer, device)
    val_loss, val_acc = evaluate(model, val_loader, criterion, device)

    if val_loss < best_val:        # improved!
        best_val = val_loss
        torch.save(model.state_dict(), "best_model.pt")   # keep only the best
        print(f"epoch {epoch}: new best val_loss {val_loss:.4f} — saved")

Adding early stopping with patience

import torch

best_val = float("inf")
patience = 5                        # stop after 5 epochs with no improvement
epochs_no_improve = 0

for epoch in range(100):
    train_loss = train_one_epoch(model, train_loader, criterion, optimizer, device)
    val_loss, val_acc = evaluate(model, val_loader, criterion, device)

    if val_loss < best_val:
        best_val = val_loss
        epochs_no_improve = 0                              # reset the counter
        torch.save(model.state_dict(), "best_model.pt")
    else:
        epochs_no_improve += 1                             # another stagnant epoch
        if epochs_no_improve >= patience:
            print(f"Early stopping at epoch {epoch}")
            break

Restoring the best model after training

import torch

# The loop may have ended on a worse epoch — load the best one back in
model.load_state_dict(torch.load("best_model.pt", map_location=device))
model.eval()
print("Restored best model for evaluation / deployment.")

4. Common Pitfalls

Pitfall 1 — Reporting the last-epoch model as your result. Without best-model saving, you present whatever the loop ended on, which may be overfit. Always monitor validation and keep the best checkpoint; report that model.

Pitfall 2 — Monitoring the training metric for stopping. Training loss almost always keeps falling, so early stopping on it never triggers (and misses overfitting entirely). Monitor a validation metric — that's the whole point.

Pitfall 3 — Patience too small (or too large). Too small a patience stops on normal validation noise before the model has converged; too large wastes compute and overfits. A few epochs (e.g. 5–10) is a reasonable starting point; tune it to how noisy your validation curve is.

5. Further Reading & Watch List

📖 Deep Learning with PyTorch, Ch. 5 & 8 — overfitting, checkpointing the best model
🌐 Learn PyTorch — Improving a model & saving the best
🛠️ PyTorch Lightning ModelCheckpoint & EarlyStopping — how frameworks automate this (preview of Module 09)
🎓 fast.ai — early stopping and the "best" epoch

⬅️ Prev: 03 · Loading for Inference · 🏁 Module & Phase 2 complete! Back to the module index or on to Module 07 — Computer Vision.