Unit 05.06: Project step: training diagnostics report

The diagnostics from this module, collected into a record that explains not just what the model scored but how it got there.

Instrumenting the loop

Per epoch: training loss, validation loss, and the gradient norm. The gradient norm is the early warning — it spikes before the loss visibly diverges, so it often identifies the exact epoch where training went wrong.

clip_grad_norm_ is called here with an enormous limit purely to *measure* the norm without altering it:

import json

import torch
from torch import nn

torch.manual_seed(0)
X = torch.randn(500, 5)
y = (X[:, 0] * 2 - X[:, 1] + 0.5 * X[:, 2]).unsqueeze(1)
Xtr, ytr, Xva, yva = X[:400], y[:400], X[400:], y[400:]

model = nn.Sequential(nn.Linear(5, 32), nn.ReLU(), nn.Linear(32, 1))
opt = torch.optim.Adam(model.parameters(), lr=0.01)

history, best = [], {"val": float("inf"), "epoch": -1}
for epoch in range(200):
    opt.zero_grad()
    train_loss = nn.MSELoss()(model(Xtr), ytr)
    train_loss.backward()
    grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1e9).item()
    opt.step()
    with torch.no_grad():
        val_loss = nn.MSELoss()(model(Xva), yva).item()
    history.append({"epoch": epoch, "train": round(train_loss.item(), 5),
                    "val": round(val_loss, 5), "grad_norm": round(grad_norm, 4)})
    if val_loss < best["val"]:
        best = {"val": round(val_loss, 5), "epoch": epoch}

print(json.dumps({
    "seed": 0,
    "optimizer": "Adam(lr=0.01)",
    "epochs": len(history),
    "first": history[0], "last": history[-1],
    "best_validation": best,
    "epochs_after_best": len(history) - 1 - best["epoch"],
    "diverged": any(h["val"] != h["val"] for h in history),
    "verdict": "stable" if history[-1]["val"] < history[0]["val"] else "did not improve",
}, indent=2))

The summary reports first and last epochs, the best validation score and where it occurred, and epochs_after_best. If that last figure is large, you trained well past the useful point.

The diverged check tests for NaN explicitly, because a NaN loss compares false against everything and can slip through a naive minimum check.

The mistake this prevents

Recording only the final epoch. Training that diverged at epoch 150 and partially recovered by 200 looks fine in a single-row summary and is deeply suspect in the history.

Takeaway

Log every epoch, watch the gradient norm, and record how long you trained past your best score.

Mark complete