The Training Loop

The basic shape

Every PyTorch training loop has the same skeleton: for each epoch, iterate over the training loader, compute loss, backprop, step the optimizer. Periodically (usually each epoch or every N steps), evaluate on the validation set with model.eval(). Save the best checkpoint. Repeat until the validation metric stops improving.

This skeleton is so universal that PyTorch users either write it themselves once and reuse it, or use a higher-level framework (PyTorch Lightning, Hugging Face Trainer, Accelerate) that writes it for them. Both are reasonable. Writing your own is the better learning experience; using a framework is the better production choice.

Tip: Write at least one training loop by hand before adopting a framework. The frameworks save typing, but they also hide the very steps you need to debug when something goes wrong.

What goes wrong in junior training loops

Forgetting opt.zero_grad() — gradients accumulate from previous steps.
Forgetting model.eval() at validation — dropout and batchnorm misbehave, metrics are noisy.
Computing loss on the GPU but printing/logging from the CPU without .item() — invisible CUDA syncs everywhere.
Computing accuracy with argmax on probabilities and argmax on logits inconsistently — same result by luck, broken when softmax changes.
Saving the model state dict at the end of training, not at the best validation step — you ship the wrong model.

The non-negotiable additions

Beyond the skeleton, every serious training loop has: gradient clipping, learning-rate scheduling, mixed-precision (autocast + GradScaler) on GPU, periodic checkpoint saving, validation against a held-out set, and logging that you can read in your monitoring tool of choice. The next track makes each of these mechanical.

Principle: A training loop is not 'done' until it can be killed and resumed without losing progress. Every saved checkpoint should include model, optimizer, scheduler, scaler, and step number.

Code

Minimal but honest training loop·python

import torch
from torch import nn, optim

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = MyModel().to(device)
opt   = optim.AdamW(model.parameters(), lr=3e-4)
sch   = warmup_cosine(opt, warmup_steps=200, total_steps=5_000)
loss_fn = nn.CrossEntropyLoss()

best_val_acc = 0.0
step = 0

for epoch in range(10):
    model.train()
    for xb, yb in train_loader:
        xb, yb = xb.to(device), yb.to(device)
        opt.zero_grad()
        logits = model(xb)
        loss = loss_fn(logits, yb)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        opt.step(); sch.step()
        step += 1

    model.eval()
    correct, total = 0, 0
    with torch.inference_mode():
        for xb, yb in val_loader:
            xb, yb = xb.to(device), yb.to(device)
            preds = model(xb).argmax(dim=-1)
            correct += (preds == yb).sum().item()
            total   += yb.size(0)
    val_acc = correct / total
    print(f"epoch {epoch} step {step} val_acc={val_acc:.4f}")

    if val_acc > best_val_acc:
        best_val_acc = val_acc
        torch.save({
            "model": model.state_dict(),
            "optimizer": opt.state_dict(),
            "scheduler": sch.state_dict(),
            "step": step, "val_acc": val_acc,
        }, "best.pt")

The basic shape

What goes wrong in junior training loops

The non-negotiable additions

Code

External links

Exercise

Progress

Comments 0