Training a Neural Network

Training a neural network means choosing its weights so that it predicts well on images it has never seen. The optimization itself, stochastic gradient descent on a loss function with gradients from backpropagation, has its own articles. This article covers the workflow around it: splitting and preparing the data, initializing the weights, monitoring progress and deciding when to stop, choosing hyperparameters, and the checks that catch bugs early.

Definition

Training is the procedure that fits a model’s parameters to a labeled training set by repeatedly reducing the average loss on it, while a held-out validation set guides choices the optimizer cannot make (how long to train, which learning rate, which architecture) and a separate test set estimates, once, how well the finished model generalizes. One pass over the training set is an epoch; each parameter update uses one mini-batch of images.

Intuition

A network with enough parameters can memorize its training images, as a student can memorize last year’s exam answers, so training loss alone says little about the goal. The workflow is built around measuring generalization honestly: keep images the model never trains on, watch how it does on them during training, and stop or change course when the two diverge.

Formal Definition

For an image classifier $f(\mathbf{x}; \theta)$ with $K$ classes, training minimizes the empirical risk of the loss function article, usually the cross-entropy between the softmax output and the label, averaged over the $N$ training images. The training loop repeats, for each mini-batch $S$ of $B$ images:

1. Forward pass: compute the class scores $f(\mathbf{x}_i; \theta)$ for $i \in S$.
2. Loss: $L = \frac{1}{B}\sum_{i \in S} \ell\big(f(\mathbf{x}_i; \theta), y_i\big)$.
3. Backward pass: compute $\nabla_\theta L$ by backpropagation.
4. Update: take an optimizer step, $\theta \leftarrow \theta - \eta\, \nabla_\theta L$ for plain SGD, or a momentum or Adam update.

The images are reshuffled at the start of each epoch. After each epoch the loss and accuracy are computed on the training and validation sets; together they form the learning curves.

Early stopping (Goodfellow, Bengio and Courville, 2016, §7.8) keeps a copy of the parameters whenever the validation error improves and stops after a fixed number of evaluations without improvement, the patience; it returns the stored parameters, not the last ones. The authors call it probably the most commonly used form of regularization in deep learning and view it as hyperparameter selection for the number of training steps, one whose many values a single run tries out.

Properties

Splitting the data

The validation set is carved out of the training data, and the test set is used for nothing that shapes the model, hyperparameters included (Goodfellow et al., §5.3). A test score that has informed decisions is no longer an unbiased estimate, so it is computed once, at the end.

Images make the splits easy to get wrong, because pictures that are nearly the same can land on both sides. Barz and Denzler (2020) found near-duplicates in the training sets for about 3.3% of the CIFAR-10 test images and 10% of the CIFAR-100 test images; after replacing them with new images, the error of the six CNN architectures they re-evaluated rose on average by 0.41 percentage points on CIFAR-10 and 2.73 on CIFAR-100. In our view the same reasoning applies to other groupings: consecutive video frames are nearly identical, and images of one patient, scene or camera share appearance a network can exploit, so splitting by video, patient or scene rather than by image keeps the validation score from rewarding memorization. Statistics computed from data, such as normalization constants, should also come from the training set alone.

Preprocessing

Inputs are usually scaled to roughly zero mean and unit variance, which (in our reasoning) matches the variance assumptions behind the initialization below. The standard image recipe subtracts a per-channel mean and divides by a per-channel standard deviation. TorchVision’s ImageNet models, for example, rescale pixel values to $[0, 1]$ and then normalize with mean $(0.485, 0.456, 0.406)$ and standard deviation $(0.229, 0.224, 0.225)$ for the red, green and blue channels; a pretrained model expects the statistics it was trained with. Data augmentation, training on randomly flipped, cropped or color-shifted copies of each image, enlarges the effective training set.

Initialization

Weights start random, because units with identical incoming weights and inputs would receive identical updates forever; Goodfellow et al. (§8.4) call breaking this symmetry perhaps the only property of the initial point known with certainty. The scale matters too: too small and the signal shrinks layer by layer, too large and it explodes. Glorot and Bengio (2010) proposed the normalized initialization, uniform weights in $\pm\sqrt{6}/\sqrt{n_{\text{in}} + n_{\text{out}}}$ (variance $2/(n_{\text{in}} + n_{\text{out}})$), to approximately preserve the variance of activations going forward and of gradients going backward. For ReLU networks, He et al. (2015) derived zero-mean Gaussian weights with standard deviation $\sqrt{2/n}$, where $n = k^2 c$ is the number of inputs to a unit of a $k \times k$ convolution over $c$ channels, with zero biases. They reported that a 30-layer network converged from their initialization while the Glorot initialization stalled.

Learning curves, overfitting and underfitting

The gap between the training and validation curves is the main diagnostic:

• Underfitting: both losses stay high. The model is too small, the learning rate is wrong, or training is too short.
• Overfitting: training loss keeps falling while validation loss levels off or rises, as in the U-shaped validation curve of Goodfellow et al.’s Figure 7.3. More data, augmentation, regularization (for example weight decay or dropout) or early stopping are the usual remedies.
• Divergence: a loss that jumps or becomes nan usually means the learning rate is too high.

Validation loss can rise while accuracy holds, because a few confident mistakes cost a lot under cross-entropy.

Hyperparameters and schedules

Goodfellow et al. (§11.4.1) single out the learning rate: if only one hyperparameter can be tuned, tune it. It is searched on a logarithmic scale, for example $10^{-5}$ to $10^{-1}$ (§11.4.4). For several hyperparameters, Bergstra and Bengio (2012) argued that random search beats a grid: for most datasets only a few hyperparameters matter, and which ones varies, so a grid wastes trials repeating the same few values of the important ones. Compared with grid-searched neural networks, random search found models as good or better within a small fraction of the computation. The learning rate is also changed during training by a learning-rate schedule, often after a short warmup; the SGD article covers step decay and warmup.

Debugging checks

Two cheap checks are common practice before a serious experiment:

• Check the loss at initialization. A classifier that starts with near-uniform softmax probabilities has cross-entropy $-\ln(1/K) = \ln K$, about $2.303$ for $K = 10$ classes. A much larger value means the network starts confidently wrong, often because the output layer is initialized too large; a much smaller one suggests a bug in the loss.
• Overfit a tiny dataset. A correct model and training loop should drive the loss on a few dozen images close to zero. Goodfellow et al. (§11.5) recommend fitting a tiny dataset to separate genuine underfitting from software defects.

Reproducibility

Training is random through initialization, shuffling, augmentation and dropout, so the seeds are recorded with the results. Seeding is necessary but not sufficient: PyTorch’s documentation states that fully reproducible results are not guaranteed across releases, platforms, or between CPU and GPU, even with identical seeds. On one machine, torch.manual_seed seeds every device, torch.use_deterministic_algorithms(True) selects deterministic implementations (and raises an error for operations that have none), and torch.backends.cudnn.benchmark = False makes cuDNN’s algorithm choice deterministic. The documentation warns that deterministic operations are often slower. Because single runs vary, settings are best compared over several seeds.

Compute and memory

The backward pass needs the activations of the forward pass, so in our reasoning activation memory grows roughly in proportion to the batch size and the number of pixels; in practice the batch size is often set by what fits in GPU memory. Mixed-precision training (Micikevicius et al., 2018) stores weights, activations and gradients in 16-bit floating point, keeps a 32-bit master copy of the weights for the updates, and multiplies the loss by a scale factor so that small gradients do not underflow to zero; the authors reported that it nearly halves memory use. In PyTorch this is torch.autocast together with torch.amp.GradScaler.

Examples

The script trains a small convolutional classifier on scikit-learn’s 1,797 grayscale 8×8 digit images, with the whole workflow on a small scale: a stratified 60/20/20 split, normalization with training-set statistics, He initialization (with a small output layer, so the initial softmax is near uniform), SGD with momentum, per-epoch learning curves, early stopping on validation loss with patience 5, and one final test evaluation. A second run trains on only 50 images to show overfitting. It runs in about three seconds on a CPU, with PyTorch 2.14.1.

import copy, math
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split

SEED = 0
torch.manual_seed(SEED)

# 1. Data: 1797 grayscale 8x8 digit images, split 60/20/20 with stratification.
digits = load_digits()
X, y = digits.images.astype(np.float32), digits.target
X_tr, X_tmp, y_tr, y_tmp = train_test_split(X, y, test_size=0.4, random_state=SEED, stratify=y)
X_va, X_te, y_va, y_te = train_test_split(X_tmp, y_tmp, test_size=0.5, random_state=SEED, stratify=y_tmp)

# 2. Normalize with statistics of the training images only (one channel here).
mean, std = X_tr.mean(), X_tr.std()
def prep(a):
    return torch.tensor((a - mean) / std).unsqueeze(1)   # shape (N, 1, 8, 8)
X_tr, X_va, X_te = prep(X_tr), prep(X_va), prep(X_te)
y_tr, y_va, y_te = (torch.tensor(t) for t in (y_tr, y_va, y_te))
print(f"train {len(y_tr)}, val {len(y_va)}, test {len(y_te)} images; "
      f"pixel mean {mean:.2f}, std {std:.2f}")

# 3. A tiny CNN with He initialization for the ReLU layers.
def make_model():
    model = nn.Sequential(
        nn.Conv2d(1, 16, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),   # 16 x 4 x 4
        nn.Conv2d(16, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),  # 32 x 2 x 2
        nn.Flatten(), nn.Linear(128, 10))
    for m in model:
        if isinstance(m, (nn.Conv2d, nn.Linear)):
            nn.init.kaiming_normal_(m.weight, nonlinearity="relu")
            nn.init.zeros_(m.bias)
    nn.init.normal_(model[-1].weight, std=0.01)   # small output layer: near-uniform softmax
    return model

@torch.no_grad()
def evaluate(model, X, y):
    model.eval()
    logits = model(X)
    return F.cross_entropy(logits, y).item(), (logits.argmax(1) == y).float().mean().item()

def train(X, y, epochs, patience=None, lr=0.01, B=32, log=()):
    model = make_model()
    opt = torch.optim.SGD(model.parameters(), lr=lr, momentum=0.9)
    loss0, _ = evaluate(model, X, y)
    print(f"loss at initialization {loss0:.3f}  (ln 10 = {math.log(10):.3f})")
    best, best_state, bad = float("inf"), None, 0
    for epoch in range(1, epochs + 1):
        model.train()
        for idx in torch.randperm(len(y)).split(B):        # reshuffle every epoch
            loss = F.cross_entropy(model(X[idx]), y[idx])  # forward + loss
            opt.zero_grad()
            loss.backward()                                # backward
            opt.step()                                     # update
        tr_loss, tr_acc = evaluate(model, X, y)
        va_loss, va_acc = evaluate(model, X_va, y_va)
        if epoch in log:
            print(f"epoch {epoch:3d}  train loss {tr_loss:.3f} acc {tr_acc:.3f}  "
                  f"val loss {va_loss:.3f} acc {va_acc:.3f}")
        if va_loss < best:                                 # early stopping bookkeeping
            best, best_state, best_epoch, bad = va_loss, copy.deepcopy(model.state_dict()), epoch, 0
        else:
            bad += 1
            if patience is not None and bad >= patience:
                break
    print(f"last epoch {epoch}; best val loss {best:.3f} at epoch {best_epoch}")
    model.load_state_dict(best_state)                      # restore the best weights
    return model

print(f"torch {torch.__version__}\n\n== full training set, early stopping (patience 5)")
model = train(X_tr, y_tr, epochs=100, patience=5, log=range(5, 101, 5))
te_loss, te_acc = evaluate(model, X_te, y_te)
print(f"test loss {te_loss:.3f}, test accuracy {te_acc:.3f}  (evaluated once)")

print("\n== 50 training images (5 per class), no early stopping")
small = torch.cat([torch.nonzero(y_tr == k)[:5, 0] for k in range(10)])
train(X_tr[small], y_tr[small], epochs=200, B=10, log=(1, 3, 10, 30, 100, 200))

Output:

train 1078, val 359, test 360 images; pixel mean 4.88, std 6.01
torch 2.14.1

== full training set, early stopping (patience 5)
loss at initialization 2.318  (ln 10 = 2.303)
epoch   5  train loss 0.065 acc 0.983  val loss 0.111 acc 0.972
epoch  10  train loss 0.018 acc 0.999  val loss 0.067 acc 0.986
epoch  15  train loss 0.008 acc 1.000  val loss 0.050 acc 0.989
epoch  20  train loss 0.005 acc 1.000  val loss 0.051 acc 0.989
last epoch 20; best val loss 0.050 at epoch 15
test loss 0.055, test accuracy 0.983  (evaluated once)

== 50 training images (5 per class), no early stopping
loss at initialization 2.320  (ln 10 = 2.303)
epoch   1  train loss 2.171 acc 0.220  val loss 2.252 acc 0.192
epoch   3  train loss 1.386 acc 0.540  val loss 1.667 acc 0.418
epoch  10  train loss 0.071 acc 1.000  val loss 0.785 acc 0.721
epoch  30  train loss 0.006 acc 1.000  val loss 0.510 acc 0.822
epoch 100  train loss 0.002 acc 1.000  val loss 0.509 acc 0.827
epoch 200  train loss 0.001 acc 1.000  val loss 0.513 acc 0.827
last epoch 200; best val loss 0.506 at epoch 49

The loss at initialization, 2.318, is close to $\ln 10 = 2.303$. Without the small output layer, using He initialization there too, the initial loss over five seeds ranged from 2.96 to 4.72: the network started out confidently wrong.

On the full training set, training accuracy reaches 100% by epoch 15 while validation loss bottoms out at 0.050; it does not improve over the next five epochs, early stopping halts at epoch 20, and the epoch-15 weights score 98.3% on the 360 test images. Over eight seeds, test accuracy ranged from 97.2% to 98.9% (median 98.3%), and stopping epochs varied from 19 to 92.

With 50 training images the model fits them perfectly by epoch 10 (the tiny-dataset check passing). Its training loss then falls to 0.001 while validation loss levels off near 0.51 and validation accuracy near 83%, against 98.9% for the full-data model: a large generalization gap that further training does not close. The output is identical across reruns on this CPU, which the PyTorch documentation does not promise across platforms.

Common Misconceptions

• “Lower training loss means a better model.” Only validation performance measures generalization; in the 50-image run training loss kept falling while validation performance did not improve.
• “Tuning on the test set is harmless if done a little.” Every decision informed by the test set biases it. Use the validation set, or cross-validation when data are scarce.
• “Setting the seed makes training reproducible.” It makes it repeatable on one software and hardware configuration at best; GPU kernels can be nondeterministic unless deterministic algorithms are requested.

Where It Is Used

Every learned vision model goes through this workflow, from small fully connected networks to convolutional networks and transformers for classification, detection, segmentation and depth estimation. Its optimization pieces are cross-entropy and other losses, backpropagation, and gradient descent with its stochastic variants and Adam; regularization, dropout and batch normalization change how a network trains and generalizes.