The training loop

Training a digit classifier

Now for a real dataset. The digits dataset holds 1,797 small images of handwritten digits. Each image is 8 × 8 pixels, so it's a list of 64 numbers. Each has a label from 0 to 9, so there are 10 classes.

The network has 64 inputs, one hidden layer of 64 neurons, and 10 outputs, one per digit.

import torch
import torch.nn as nn
import matplotlib.pyplot as plt
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split

digits = load_digits()
X = torch.tensor(digits.data / 16.0, dtype=torch.float32)     # pixels scaled to 0–1
y = torch.tensor(digits.target)
X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2, random_state=0)

torch.manual_seed(0)
model = nn.Sequential(nn.Linear(64, 64), nn.ReLU(), nn.Linear(64, 10))
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.5)

history = {"train": [], "val": [], "acc": []}
for epoch in range(150):
    model.train()
    optimizer.zero_grad()
    loss = loss_fn(model(X_train), y_train)
    loss.backward()
    optimizer.step()

    model.eval()
    with torch.no_grad():
        logits = model(X_val)
        history["train"].append(loss.item())
        history["val"].append(loss_fn(logits, y_val).item())
        history["acc"].append((logits.argmax(dim=1) == y_val).float().mean().item())
    if epoch % 30 == 0 or epoch == 149:
        print(f"epoch {epoch:>3}  train loss {history['train'][-1]:.3f}  val loss {history['val'][-1]:.3f}  val accuracy {history['acc'][-1]:.3f}")

plt.plot(history["train"], label="training loss")
plt.plot(history["val"], label="validation loss")
plt.xlabel("epoch")
plt.legend()

Here's what the new lines do, from the top:

• The pixel values run from 0 to 16. Dividing by 16 squeezes them into 0 to 1.
• train_test_split holds back 20% of the images (test_size=0.2) as the validation set. random_state=0 makes the split the same every run.
• torch.manual_seed(0) fixes the model's random starting weights.
• nn.CrossEntropyLoss() is the cross-entropy loss from Phase 2. It applies softmax for you.
• The history dictionary keeps three lists: training loss, validation loss and validation accuracy. .item() turns a one-number tensor into a plain Python number.
• The code prints a line every 30 epochs, plus the last one. The plot draws both loss curves.

The accuracy line needs a closer look. logits.argmax(dim=1) picks the position of the highest score in each row. That position is the model's guess. Comparing it with y_val gives True or False for each image. .float() turns those into 1s and 0s, and .mean() averages them. The result is the fraction of guesses that were right. For example, 3 right out of 4 gives 0.75.

In this lesson, each epoch is a single step on all the training data at once. The next lesson splits each epoch into many smaller steps.

Evaluation mode, and why it's two calls

Think of a sports team. In practice drills, the coach benches random players so nobody leans on one star. On game day, everyone plays. Dropout works like those drills, and model.eval() is game day.

You need two separate switches when you evaluate. They do different jobs:

• model.eval() switches some layers to their testing behavior. Dropout stops switching off neurons. Batch normalization uses the averages it stored during training. It does not turn off gradients.
• torch.no_grad() stops PyTorch from recording steps for a backward pass. It does not change how any layer behaves.

So use both when you evaluate. Then call model.train() again before the next training step.

Forgetting either switch causes quiet problems. Forget eval() and a model with dropout gives noisy validation numbers that make it look worse than it really is. Forget train() afterwards and dropout stays off for the rest of training.

Keep the best, not the last

A runner remembers their personal best, not their latest time. Training should work the same way.

Validation loss usually falls for a while. Then it reaches a low point and starts to creep back up. That's the model starting to overfit: it learns quirks of the training data that don't carry over to new data.

So save a copy of the weights each time the validation loss hits a new low. A saved copy like this is called a checkpoint. At the end, load the best one back. You can also stop training once the validation loss hasn't improved for a while. That's called early stopping.

best_val, best_state = float("inf"), None
for epoch in range(num_epochs):
    ...                                              # train and compute val_loss
    if val_loss < best_val:
        best_val = val_loss
        best_state = {k: v.clone() for k, v in model.state_dict().items()}
model.load_state_dict(best_state)                    # back to the best checkpoint
torch.save(best_state, "digits.pt")                  # on your machine: save to disk

Here's how it works:

• float("inf") is infinity. Any real loss beats it, so the first epoch always counts as the best so far.
• model.state_dict() is a dictionary of every weight tensor in the model, by name.
• .clone() makes a separate copy of each tensor.
• model.load_state_dict(best_state) puts the saved weights back into the model.
• torch.save writes them to a file. That works on your own machine, not in the browser.

Reading the curves

The two loss curves tell you how training is going. This table lists the common patterns:

"nan" means "not a number". It's what you get when the numbers grow too big and blow up.

The third row has a handy test. Before any training, a model with 10 classes is guessing. Its loss should start near 2.30, which is ln 10. If your first loss is far from that, check your code.

import math
print(round(math.log(10), 3))   # 10 classes
print(round(math.log(2), 3))    # 2 classes

This is the same diagnosis as Phase 3's complexity curves, now drawn over training time.