The training loop, step by step

The idea

Unlike scikit-learn's single fit() call, PyTorch hands you the training loop to write yourself. That's more code, and also full control: you decide how batches are drawn, what gets logged, when to evaluate and when to stop. Every PyTorch project has a loop with the same skeleton, so once you know it, you can read anyone's training code.

for each epoch (a full pass over the training data):
    model.train()
    for each mini-batch:
        zero gradients → forward → loss → backward → optimizer step
    model.eval()
    with no gradients: compute validation loss
    log both losses

The data

Coffee demand again: 4 features, predict units sold. Standardise inputs (gradient-based training needs comparable scales), and keep the last 66 days for validation:

import numpy as np
import pandas as pd
import torch
from torch import nn

demand = pd.read_csv("daily_demand.csv")
FEATURES = ["temperature", "weekend", "promo", "ad_spend"]
X = demand[FEATURES].to_numpy(dtype=np.float32)
y = demand["units"].to_numpy(dtype=np.float32).reshape(-1, 1)

mean, std = X[:300].mean(axis=0), X[:300].std(axis=0)     # statistics from training rows only
X = (X - mean) / std
X_train, y_train = torch.from_numpy(X[:300]), torch.from_numpy(y[:300])
X_val, y_val = torch.from_numpy(X[300:]), torch.from_numpy(y[300:])
print(X_train.shape, y_train.shape, X_val.shape)

torch.Size([300, 4]) torch.Size([300, 1]) torch.Size([66, 4])

The loop

Here's a complete training loop for the demand data: a small network, an optimizer, a loss function, and 200 passes over the training data in batches of 32. Every PyTorch project you'll see has a loop shaped like this one:

torch.manual_seed(0)
model = nn.Sequential(nn.Linear(4, 16), nn.ReLU(), nn.Linear(16, 1))
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
loss_fn = nn.MSELoss()
batch_size = 32

for epoch in range(1, 201):
    model.train()
    order = torch.randperm(len(X_train))              # shuffle every epoch
    for start in range(0, len(X_train), batch_size):
        idx = order[start:start + batch_size]
        optimizer.zero_grad()
        loss = loss_fn(model(X_train[idx]), y_train[idx])
        loss.backward()
        optimizer.step()

    model.eval()
    with torch.no_grad():
        train_loss = loss_fn(model(X_train), y_train).item()
        val_loss = loss_fn(model(X_val), y_val).item()
    if epoch in (1, 10, 50, 100, 200):
        print(f"epoch {epoch:3d}  train MSE {train_loss:8.2f}  val MSE {val_loss:8.2f}")

epoch   1  train MSE  3105.28  val MSE  3449.84
epoch  10  train MSE   863.41  val MSE   641.93
epoch  50  train MSE    28.19  val MSE    36.46
epoch 100  train MSE    24.25  val MSE    29.97
epoch 200  train MSE    23.10  val MSE    28.22

Line by line:

• Epochs: each epoch is one pass over all 300 training rows.
• torch.randperm shuffles the row order each epoch, so batches differ every time (the "stochastic" in stochastic gradient descent).
• Mini-batches of 32 rows: about 10 updates per epoch. Each update uses a noisy estimate of the full gradient: cheap and effective.
• model.train() / model.eval() switch layers that behave differently in training (dropout, batch normalisation; module 12). Calling them always is a good habit, even when your model has no such layers yet.
• torch.no_grad() for validation: no graph recording, less memory.
• Logging both losses is how you watch learning happen and spot overfitting.

The MSE drops from thousands (random weights predicting nonsense) to around 25, in line with the linear regression from the ML track (RMSE ≈ 5 units). On this nearly linear data, the network matches linear regression; it can't magically beat the noise.

Reading the curves

Making predictions

After training, predict in eval mode without gradients, and convert back to NumPy:

model.eval()
with torch.no_grad():
    preds = model(X_val).squeeze().numpy()
mae = np.mean(np.abs(preds - y_val.squeeze().numpy()))
print(preds[:5].round(1), round(float(mae), 2))

[76.9 46.1 67.9 47.  53.8] 4.39

The validation MAE is comparable to the linear regression in the ML track: a good sanity check that the loop works.

Why this matters for clean code

Training loops grow features over time: logging, checkpoints, early stopping, metrics. Start with the clean skeleton above, and when it grows, extract train_one_epoch(model, loader, optimizer, loss_fn) and evaluate(model, loader, loss_fn) functions. Each then fits on one screen and can be tested with a tiny dataset. The project-structure module builds exactly that.

Recap

• Epoch = one pass over the training data; mini-batches give several updates per epoch.
• Shuffle each epoch (torch.randperm or a DataLoader).
• Per batch: optimizer.zero_grad(), forward, loss, backward(), optimizer.step().
• model.train() for training, model.eval() + torch.no_grad() for validation and prediction.
• Log training and validation loss every epoch and read the curves for overfitting or bugs.