The Training Loop

You know what a neural network computes, and you know autograd can produce gradients for free. This lesson connects the two: the training loop, the handful of lines at the heart of every PyTorch project. Once you can read and write this loop, everything from linear regression to giant language models is a variation on the same theme.

Three ingredients before you train: MCO

Every training script starts by preparing three objects. A handy mnemonic: MCO — Model, Criterion, Optimizer.

M is for Model

The thing that makes predictions. PyTorch's nn module gives you ready-made layers; nn.Sequential chains them into a network:

import torch
from torch import nn
 
# 4 inputs -> hidden(3, ReLU) -> hidden(4, ReLU) -> 3 outputs
model = nn.Sequential(
    nn.Linear(4, 3),
    nn.ReLU(),
    nn.Linear(3, 4),
    nn.ReLU(),
    nn.Linear(4, 3),
)
 
print(model)
print(model.state_dict().keys())   # every learnable weight and bias

Each nn.Linear(in, out) is exactly the weighted-sum-plus-bias you built in NumPy — a weight matrix and a bias vector, initialized randomly. The state_dict() is the model's memory: all its learnable parameters, which is what training will change and what you'll later save to disk.

C is for Criterion

The criterion (loss function) measures how wrong the predictions are — one number, where lower is better. You pick it based on the task:

criterion = nn.MSELoss()              # regression (model ends in plain Linear)
criterion = nn.BCEWithLogitsLoss()    # binary classification (1 output, raw logit)
criterion = nn.CrossEntropyLoss()     # multiclass (N outputs, raw logits)

Note a modern PyTorch convention: for classification the model outputs raw logits — no sigmoid or softmax layer at the end. BCEWithLogitsLoss and CrossEntropyLoss apply the squashing internally, which is both numerically more stable and less code. Only add a sigmoid/softmax yourself when you need actual probabilities at prediction time.

O is for Optimizer

The optimizer updates the weights using the gradients. It's gradient descent with an engine attached:

from torch import optim
 
optimizer = optim.SGD(model.parameters(), lr=0.01)    # classic gradient descent
optimizer = optim.Adam(model.parameters(), lr=0.001)  # adaptive — great default

You hand it model.parameters() — the list of tensors it's allowed to modify — and a learning rate. You already know why the learning rate is the most important knob; here's a refresher on the three regimes:

Too small and training crawls; too large and the loss bounces or explodes. Adam adapts a per-parameter step size on the fly, which makes it far more forgiving of the initial learning-rate choice — that's why Adam (or its sibling AdamW) with lr=0.001 is the standard starting point, while plain SGD typically needs more tuning to shine.

The loop: forward, loss, backward, step

With MCO in place, one training step is four moves:

1. Forward pass — output = model(x): push data through the network.
2. Compute loss — loss = criterion(output, y): one number measuring error.
3. Backward pass — loss.backward(): autograd fills .grad for every parameter.
4. Update — optimizer.step(): nudge every weight downhill, w ← w − lr·grad.

And one bookkeeping move: optimizer.zero_grad(). Remember from the last lesson that gradients accumulate — each backward() adds to .grad instead of overwriting it. Without zeroing, step 2's gradients would stack on top of step 1's, and the updates would be garbage. So every iteration clears the gradients before (or right after) the update.

Repeat the whole dance many times. Each full pass through the training data is called an epoch.

A complete minimal example

Let's train the smallest possible network — a single nn.Linear(1, 1), i.e. y = wx + b — to recover a line from noisy synthetic data. Every real training script you'll ever write has this exact skeleton:

import torch
from torch import nn, optim
import matplotlib.pyplot as plt
 
torch.manual_seed(42)
 
# Synthetic data: y = 2x - 1 + noise
X = torch.rand(100, 1) * 10                 # shape [100, 1]
y = 2 * X - 1 + torch.randn(100, 1) * 0.8   # shape [100, 1]
 
# MCO
model = nn.Linear(1, 1)
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.1)
 
# Training loop
losses = []
for epoch in range(200):
    output = model(X)                # 1. forward
    loss = criterion(output, y)      # 2. loss
 
    optimizer.zero_grad()            #    clear old gradients
    loss.backward()                  # 3. backward
    optimizer.step()                 # 4. update
 
    losses.append(loss.item())
    if (epoch + 1) % 50 == 0:
        print(f"epoch {epoch+1:3d} | loss {loss.item():.4f}")
 
w, b = model.weight.item(), model.bias.item()
print(f"\nlearned: y = {w:.2f}x + {b:.2f}   (true: y = 2.00x - 1.00)")
 
plt.plot(losses)
plt.xlabel("epoch"); plt.ylabel("MSE loss")
plt.show()

That's it. A 175-million-parameter network trains with the same eight lines — only the model, the data, and the number of epochs change.

Reading the loss curve

The loss-per-epoch plot is your training EKG. Learn to read these shapes — they'll tell you what to fix (plotting a validation loss alongside the training loss, which we'll set up properly in the DataLoader lesson, makes the diagnosis even sharper):

• Smooth decline that flattens out at a low value — healthy. Training converged.
• Still clearly falling when the loop ends — underfitting by impatience: train longer, or raise the learning rate a bit.
• Flattens early at a high value — underfitting by capacity: the model is too simple for the pattern, or the learning rate is too small to make progress.
• Training loss keeps dropping while validation loss turns around and rises — overfitting: the model has started memorizing the training set. More data, regularization (dropout, coming soon), or early stopping.
• Spiky, oscillating, or growing — learning rate too high. The optimizer is overshooting the valley. Cut lr by 10x and try again.

Sensible defaults, in one place

When in doubt, start here and adjust only when the loss curve tells you to:

Two pairings to burn in: CrossEntropyLoss wants raw logits and integer class labels (torch.long, not one-hot). MSELoss wants the target shaped exactly like the output — a [100] target against a [100, 1] output will "work" via broadcasting and silently ruin your training.

Next up: nn.Module — the class-based way to define networks, plus just enough object-oriented Python to read any PyTorch model ever written.