Module 5 — Hand-written training and evaluation loop

The four previous modules built the pieces: tensors, autograd, an nn.Module, a DataLoader. This one assembles them into the smallest complete training program that actually learns. It is deliberately hand-written — no fit, no trainer library — because owning the five canonical steps of a PyTorch loop makes every framework built on top of them, and every bug encountered in the wild, immediately readable.

At the end of this module you have the first working version of the Fashion-MNIST classifier: it trains for a few epochs on a laptop CPU, reports loss and accuracy per epoch, and cleanly separates training from evaluation. From here on, modules 6 to 10 add capabilities without ever changing the shape of the loop.

The five canonical steps

Every PyTorch training iteration performs, in this exact order:

optimizer.zero_grad()              # 1. clear gradients
outputs = model(inputs)            # 2. forward
loss = criterion(outputs, targets) # 3. compute loss
loss.backward()                    # 4. backward
optimizer.step()                   # 5. update weights

The order is not decorative. zero_grad first, because gradients accumulate as we saw in module 2. Forward next, to build the graph on the current inputs. Loss computed on the outputs. backward walks that graph to fill .grad. step reads .grad and updates the weights. Any reordering either loses gradients, applies stale updates, or breaks silently.

A complete loop, from scratch

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

# --- data (module 4) ---
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.2860,), (0.3530,)),
])
train_set = datasets.FashionMNIST("data", train=True, download=True, transform=transform)
val_set   = datasets.FashionMNIST("data", train=False, download=True, transform=transform)
train_loader = DataLoader(train_set, batch_size=64, shuffle=True, num_workers=2)
val_loader   = DataLoader(val_set,   batch_size=256, shuffle=False, num_workers=2)

# --- model (module 3) ---
class FashionMLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Flatten(),
            nn.Linear(784, 256), nn.ReLU(), nn.Dropout(0.3),
            nn.Linear(256, 256), nn.ReLU(),
            nn.Linear(256, 10),
        )
    def forward(self, x):
        return self.net(x)

model = FashionMLP()
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

# --- training ---
for epoch in range(10):
    model.train()
    running_loss, running_correct, seen = 0.0, 0, 0
    for x, y in train_loader:
        optimizer.zero_grad()
        logits = model(x)
        loss = criterion(logits, y)
        loss.backward()
        optimizer.step()

        running_loss += loss.item() * y.size(0)
        running_correct += (logits.argmax(1) == y).sum().item()
        seen += y.size(0)

    train_loss = running_loss / seen
    train_acc  = running_correct / seen

    # --- evaluation ---
    model.eval()
    val_loss, val_correct, val_seen = 0.0, 0, 0
    with torch.no_grad():
        for x, y in val_loader:
            logits = model(x)
            val_loss += criterion(logits, y).item() * y.size(0)
            val_correct += (logits.argmax(1) == y).sum().item()
            val_seen += y.size(0)

    print(f"epoch {epoch:02d}  "
          f"train loss {train_loss:.4f}  train acc {train_acc:.4f}  "
          f"val loss {val_loss/val_seen:.4f}  val acc {val_correct/val_seen:.4f}")

Read that block twice: it is the reference shape every later module refines. Ten epochs on a laptop CPU take a couple of minutes and reach around 89 % validation accuracy — a decent baseline for an MLP on Fashion-MNIST, and a starting point we will improve.

CrossEntropyLoss expects logits, not probabilities

The model outputs raw scores. It does not end with a softmax. nn.CrossEntropyLoss combines a log_softmax and a negative log-likelihood in one numerically stable operation. Applying softmax yourself before feeding into CrossEntropyLoss runs without error but shifts the entire loss landscape and slows training drastically. This is one of the two most frequent bugs when transitioning from TensorFlow, where SparseCategoricalCrossentropy(from_logits=True) requires the same convention explicitly.

Targets for classification are class indices, of dtype int64. A tensor of one-hot labels would fail with a clear type error. The Fashion-MNIST dataset already returns integer labels, so no conversion is needed.

model.train() versus model.eval(): not a decoration

model.train() and model.eval() do exactly two things, and they matter enormously.

Dropout layers zero out random activations during training and pass everything through during evaluation. BatchNorm2d computes statistics from the batch during training and uses stored running statistics during evaluation. Both switches are controlled by whether the module is in train mode or eval mode.

Forgetting model.eval() before validation lets dropout run on validation inputs — a fraction of activations gets randomly zeroed — and lets BatchNorm compute statistics from tiny validation batches, further poisoning the running averages. The reported validation accuracy is worse than the real one, sometimes by several percent. Symmetrically, forgetting model.train() at the start of the next epoch leaves the model in eval mode and disables both regularisers.

eval() and no_grad() are complementary, not alternatives

model.eval() changes layer behaviour. torch.no_grad() changes whether gradients are recorded. Evaluation code needs both. Using only no_grad() gives correct gradients (none) but wrong dropout and BatchNorm; using only eval() computes the right forward pass but wastes memory and time building an unused graph.

Reading the two curves side by side

The four numbers printed per epoch — train loss, train accuracy, validation loss, validation accuracy — encode the entire training story if you know how to read them.

Pattern	Diagnosis
Both losses fall together, gap stays small	Healthy training; keep going
Train loss falls, val loss rises for several epochs	Overfitting starts; regularise or stop earlier
Both losses stall well above zero	Underfitting; larger model or longer training
Val loss jumps epoch to epoch	Learning rate too high, or a data leak between batches
Train loss falls, val stays flat from the start	Distribution shift or bug in the validation pipeline

Overfitting is the most common of these on a small model like our MLP. The classical response is not more training but more regularisation: raise dropout, add weight decay (module 6), or reduce the network's capacity. Training longer just moves the training loss further below the validation loss without changing the story.

Logging what matters

The loss.item() * y.size(0) weighting looks strange until you realise batches can have different sizes. loss is already the mean over the batch, so multiplying by the batch size gives the sum, and dividing by the total examples at the end gives the true epoch mean. Reporting just the last batch's loss, as some tutorials do, is misleading whenever the last batch is small or unusually easy.

(logits.argmax(1) == y).sum().item() is the accuracy count. .item() extracts a Python integer from a zero-dim tensor. Chaining .detach().cpu().item() matters when the tensor lives on a GPU; on the CPU, .item() alone is fine.

Save the four numbers per epoch to a CSV

A three-line CSV of epoch, train loss, val loss, val accuracy is worth more than any dashboard for the first iteration of a project. It survives kernel restarts, can be diffed across experiments, and plots in one command. Add the timestamp and you have training run reproducibility for free.

In summary

• Every PyTorch iteration follows five ordered steps: zero_grad, forward, loss, backward, step; changing the order or removing one silently corrupts training.
• nn.CrossEntropyLoss expects logits, not probabilities, and integer class indices; applying softmax before it is one of the most common bugs.
• model.train() and model.eval() control Dropout and BatchNorm behaviour; validation needs both model.eval() and torch.no_grad(), and they cannot substitute for each other.
• Read train and validation curves side by side: a growing gap signals overfitting, a stalled pair signals underfitting, and a jumpy validation loss signals a learning rate that is too high.

Next module: optimisers and learning-rate schedules, which turn this baseline into a properly tuned classifier.