48 - A neural network from scratch in NumPy, then in PyTorch

You have all the pieces: the neuron, the layers, the forward pass, backpropagation, gradient descent. It is like having the engine, the wheels and the chassis on the table: the car is missing. This lesson assembles it. You will build a complete training loop, the program that takes real data, splits it honestly, trains a network on mini-batches, monitors its progress, stops it at the right moment and evaluates it just once on data it has never seen. First you will do it entirely by hand with NumPy, with three optimizers written from scratch, then with a library (scikit-learn and PyTorch, both run here), recognizing in every line something you have already written. This is the lesson where everyday work with networks stops being theory.

The problem is handwritten digit recognition: scikit-learn's Digits dataset contains 1797 grayscale images of 8×8 pixels, so 64 features per example, with 10 classes (the digits 0 to 9). It is not the famous MNIST (28×28 pixels, 70 000 images), but it has the same structure and runs in a few seconds on a laptop, which lets you run dozens of experiments (the dataset is described in the scikit-learn documentation, listed in the sources).

The vocabulary of the training loop

• Iteration (or step): a single weight update, computed on a mini-batch.
• Epoch: a complete pass over all the training examples. With N examples and a mini-batch of size B, an epoch has N/B iterations.
• Mini-batch: the subset of examples used to estimate the gradient at each iteration (gradient descent lesson).
• Training, validation, test: the three sets from the first lesson of the course. Here we will use them precisely: training serves to update the weights, validation to monitor overfitting and decide when to stop, test to give the final number, once only.
• Learning curve: the plot (or table) of loss and accuracy on training and validation, epoch by epoch. It is the main diagnostic tool.
• Early stopping: stopping training when the validation loss stops improving, and restoring the best weights.

The six entries are the anchor points for the rest of the lesson, and they must be kept distinct: the iteration is the unit of computation, the epoch is the unit of time in which you read the curve, and training, validation and test are three different roles of the data, not three quantities. The useful relation to remember is N/B: with the 1077 training examples and mini-batches of 32, an epoch is 34 iterations, and this number comes back in the experiment on batch size.

The diagram shows the loop. For every mini-batch the five steps are carried out; at the end of the epoch you evaluate on the validation set and decide whether to continue. Everything that follows is the realization, line by line, of this scheme.

Preparing the data: splitting and standardization, in the right order

The first part is the most important, and the one where the most silent mistakes are made. We split into three sets in a stratified way (the same proportion of digits in each) and before any transformation. The standardization learns mean and standard deviation from the training set alone, and applies them to validation and test.

Note on the code: the Python blocks in this article must be run in order, in a single script or notebook. Each one reuses the imports and variables defined in the previous blocks.

import numpy as np
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

X, y = load_digits(return_X_y=True)
X_tmp, X_te, y_tmp, y_te = train_test_split(X, y, test_size=0.2, stratify=y, random_state=0)          # 20% test
X_tr, X_va, y_tr, y_va = train_test_split(X_tmp, y_tmp, test_size=0.25, stratify=y_tmp, random_state=0)  # 60% train, 20% val

scaler = StandardScaler().fit(X_tr)                          # parameters learned from the training set ONLY
X_tr, X_va, X_te = scaler.transform(X_tr), scaler.transform(X_va), scaler.transform(X_te)

print("shapes  train %s  validation %s  test %s" % (X_tr.shape, X_va.shape, X_te.shape))
print("share of each digit in the training set (stratified):", np.round(np.bincount(y_tr) / len(y_tr), 3))

shapes  train (1077, 64)  validation (360, 64)  test (360, 64)
share of each digit in the training set (stratified): [0.099 0.101 0.098 0.101 0.101 0.101 0.101 0.099 0.097 0.1  ]

What the code does. The first train_test_split sets aside 20% for the test; the second splits the rest into training (75% of the remainder, that is 60% of the total) and validation (25%, that is 20% of the total). stratify keeps the proportions of the 10 digits in every set, and random_state makes everything reproducible. The StandardScaler is trained (fit) only on X_tr; the other two matrices are only transformed. Had you trained the scaler on the whole dataset, the test would have influenced the training: the data leakage seen in the first lesson. With 10 almost balanced classes, accuracy is a reasonable metric; with imbalanced classes you should also look at precision and recall.

The network, the optimizer and the loop, entirely by hand

We rewrite the network of the backpropagation lesson (dense layers, ReLU in the hidden ones, softmax with cross-entropy at the output) with two novelties: the weight update is entrusted to a passo method that can use three different optimizers, and there is a method to measure loss and accuracy. The three optimizers:

• Plain SGD: θ ← θ − α·g.
• SGD with momentum: the idea is a ball rolling down a hill, gathering speed. You keep a "velocity" v, a moving average of past gradients, and update v ← 0.9·v + g, θ ← θ − α·v. It crosses flat regions faster and damps the oscillations in narrow valleys.
• Adam (Kingma and Ba, 2015): it keeps two moving averages for each parameter, one of the gradient (m, like momentum) and one of its square (s). The update is θ ← θ − α·m̂/(√ŝ + ε), where m̂ and ŝ are the averages corrected for the bias of the start. In practice every parameter has its own adaptive learning rate: parameters with large, noisy gradients take smaller steps, those with small, steady gradients take larger ones.

The three optimizers are three answers to the same question, how to use the gradient g. SGD uses it as it is, and therefore inherits its noise; momentum turns it into a moving average, which damps the oscillations; Adam also normalizes the scale, parameter by parameter. In the code that follows, the difference is all in the passo method: a few lines of array arithmetic, with nothing hidden.

def softmax(Z):
    Z = Z - Z.max(axis=1, keepdims=True)
    E = np.exp(Z)
    return E / E.sum(axis=1, keepdims=True)

class Rete:
    def __init__(self, dimensioni, seed=0, init="he"):
        rng = np.random.default_rng(seed)
        coppie = list(zip(dimensioni[:-1], dimensioni[1:]))
        self.W = [rng.normal(size=(i, o)) * np.sqrt(2.0 / i) if init == "he" else np.zeros((i, o)) for i, o in coppie]
        self.b = [np.zeros(o) for _, o in coppie]
        self.stato = {}                                           # optimizer memory (velocities, moving averages)

    def avanti(self, X):
        self.Z, self.A = [], [X]
        A = X
        for l, (W, b) in enumerate(zip(self.W, self.b)):
            Z = A @ W + b
            A = softmax(Z) if l == len(self.W) - 1 else np.maximum(0, Z)
            self.Z.append(Z); self.A.append(A)
        return A

    def indietro(self, y):
        m = len(y)
        delta = self.A[-1] - np.eye(self.A[-1].shape[1])[y]
        gW, gb = [None] * len(self.W), [None] * len(self.W)
        for l in range(len(self.W) - 1, -1, -1):
            gW[l] = self.A[l].T @ delta / m
            gb[l] = delta.sum(axis=0) / m
            if l > 0:
                delta = (delta @ self.W[l].T) * (self.Z[l - 1] > 0)
        return gW, gb

    def passo(self, gW, gb, ottimizzatore, lr, t):
        for k, (p, g) in enumerate(zip(self.W + self.b, gW + gb)):       # p is an array: the update is in-place
            if ottimizzatore == "sgd":
                p -= lr * g
            elif ottimizzatore == "momentum":
                v = self.stato.setdefault(("v", k), np.zeros_like(p))
                v *= 0.9; v += g
                p -= lr * v
            else:                                                         # adam
                m1 = self.stato.setdefault(("m", k), np.zeros_like(p))
                m2 = self.stato.setdefault(("s", k), np.zeros_like(p))
                m1 *= 0.9;   m1 += 0.1 * g
                m2 *= 0.999; m2 += 0.001 * g * g
                p -= lr * (m1 / (1 - 0.9 ** t)) / (np.sqrt(m2 / (1 - 0.999 ** t)) + 1e-8)

    def perdita(self, X, y):
        P = self.avanti(X)
        return -np.log(np.clip(P[np.arange(len(y)), y], 1e-12, 1.0)).mean()

    def accuratezza(self, X, y):
        return (self.avanti(X).argmax(axis=1) == y).mean()

def addestra(rete, epoche, batch, lr, ottimizzatore, seed=0, pazienza=None):
    rng = np.random.default_rng(seed)
    t, storia, migliore = 0, [], (np.inf, None, 0)
    for epoca in range(epoche):
        ordine = rng.permutation(len(y_tr))                              # reshuffle every epoch
        for inizio in range(0, len(ordine), batch):
            idx = ordine[inizio:inizio + batch]
            rete.avanti(X_tr[idx])
            gW, gb = rete.indietro(y_tr[idx])
            t += 1
            rete.passo(gW, gb, ottimizzatore, lr, t)
        storia.append((rete.perdita(X_tr, y_tr), rete.perdita(X_va, y_va), rete.accuratezza(X_tr, y_tr), rete.accuratezza(X_va, y_va)))
        if pazienza:                                                    # early stopping
            if storia[-1][1] < migliore[0] - 1e-4:
                migliore = (storia[-1][1], [p.copy() for p in rete.W + rete.b], epoca)
            elif epoca - migliore[2] >= pazienza:
                break
    return storia, migliore

print("optimizer      lr       val. loss after 1, 5, 10, 30 epochs                final val. accuracy")
for nome, lr in (("sgd", 0.1), ("momentum", 0.05), ("adam", 0.002)):
    rete = Rete([64, 32, 10])
    storia, _ = addestra(rete, epoche=30, batch=32, lr=lr, ottimizzatore=nome)
    print("%-13s  %-6g  %-45s  %.3f" % (nome, lr, [round(float(storia[i][1]), 3) for i in (0, 4, 9, 29)], storia[-1][3]))

optimizer      lr       val. loss after 1, 5, 10, 30 epochs                final val. accuracy
sgd            0.1     [0.744, 0.285, 0.236, 0.212]                   0.958
momentum       0.05    [0.393, 0.277, 0.265, 0.268]                   0.967
adam           0.002   [1.379, 0.372, 0.262, 0.238]                   0.958

What the code does. The class contains everything you have already written (avanti, indietro with the four equations) plus the passo method. An important technical detail: p -= lr * g modifies the array in-place; since p is the same object contained in self.W or self.b, the update reaches the network. For momentum and Adam the network keeps the velocity and moving-average arrays in self.stato. The addestra function is the loop of the diagram: for every epoch it reshuffles the indices, runs through the training set in mini-batches, does avanti, indietro and passo, then records loss and accuracy on training and validation; if patience is enabled it applies early stopping.

The comparison shows different behaviors. After a single epoch momentum is already at 0.39 validation loss, plain SGD at 0.74 and Adam at 1.38 (Adam with a small learning rate starts more slowly). After 5 and 10 epochs they all get closer, and at 30 epochs they sit around 0.21-0.27 with validation accuracy between 0.958 and 0.967. There is no absolute winner on this small problem: Adam is widely used because it is less sensitive to the choice of learning rate on large and varied problems, but here plain gradient descent with momentum does just as well. If you compare optimizers, remember that each has its own suitable learning rate: using the same value for all of them would actually compare how unlucky they are with that value.

Why momentum and Adam help: an intuition with arithmetic

Imagine going down a long, narrow valley, steep on both sides and almost flat along the bottom. The gradient has a large component that bounces from one side to the other (alternating sign: +1, −1, +1, ...) and a small but constant component along the bottom (for example +0.1 at every step). With plain SGD the first makes you zigzag and the second makes you advance slowly. With momentum, the velocity is a weighted sum of past gradients: v ← 0.9·v + g.

The alternating components cancel out (the sum of +1, −1, +1, ... with decreasing weights stays small), while the constant one accumulates like a geometric series: v = g·(1 + 0.9 + 0.9² + ...) = g/(1 − 0.9) = 10·g. Along the bottom of the valley the effective step is ten times larger, while the zigzag is damped. That is why momentum accelerates where the gradient is consistent.

Adam goes one step further: it divides the update by the square root of the moving average of the squared gradient. If a parameter has a gradient that is always around 0.001, that root is about 0.001 and the ratio gradient/root is about 1: the parameter takes a step of roughly the size of the learning rate. If another has a gradient that is always around 100, the root is about 100 and the ratio is still about 1: the same step.

In practice Adam normalizes the scale of the gradients, parameter by parameter, which is why it works with very similar learning rates on very different problems, and why it uses typical values of 0.001-0.003 instead of the 0.1 of SGD. The price is that adapting to noisy gradients can, in some problems, generalize a little worse than a well-tuned SGD: there is no absolutely best optimizer.

The learning rate, measured

In the gradient descent lesson you saw the learning rate on the parabola. Let us look at it on a real network, with plain SGD and 20 epochs:

print("learning rate   final val. loss   final val. accuracy")
for lr in (0.001, 0.01, 0.1, 0.5, 2.0):
    rete = Rete([64, 32, 10])
    with np.errstate(all="ignore"):
        storia, _ = addestra(rete, epoche=20, batch=32, lr=lr, ottimizzatore="sgd")
    print("%-13g   %15.4f   %19.3f" % (lr, storia[-1][1], storia[-1][3]))

learning rate   final val. loss   final val. accuracy
0.001                    1.8588                 0.428
0.01                     0.4757                 0.897
0.1                      0.2189                 0.961
0.5                      0.2753                 0.964
2                        8.4428                 0.694

What the code does. It repeats the same training changing only the learning rate. With 0.001 the model learns too slowly: after 20 epochs the loss is 1.86 and the validation accuracy 0.43, still far from converging. With 0.01 it improves (0.476, 0.897), with 0.1 it is close to the optimum (0.219, 0.961); 0.5 is still good (0.275, 0.964) but noisier; with 2.0 the steps are too large and the loss explodes to 8.44, with accuracy 0.69. The lesson is the same as on the parabola, but now you see its concrete effect: the learning rate is the hyperparameter to settle first, and it is explored on a logarithmic scale. The 2.0 row is a diverging run, so its exact figures depend on the library versions: on a second machine with NumPy 2.4.6 the same run ended at 13.5853 and 0.508 (the other four rows were identical), while the figures above come from NumPy 2.3.5 and scikit-learn 1.8.0.

Reading the learning curves: overfitting and early stopping

We train with Adam for up to 200 epochs, but with a patience of 10: if the validation loss does not improve for 10 consecutive epochs we stop and restore the weights of the best epoch. Then, and only now, we open the test set.

rete = Rete([64, 32, 10])
storia, migliore = addestra(rete, epoche=200, batch=32, lr=0.002, ottimizzatore="adam", pazienza=10)
epoca_migliore = migliore[2]

print("epochs run: %d   best epoch: %d" % (len(storia), epoca_migliore))
print("epoch   train loss   val. loss   train acc.   val. acc.")
for e in (0, 5, 10, epoca_migliore, len(storia) - 1):
    print("%5d   %10.4f   %9.4f   %10.3f   %9.3f" % ((e,) + storia[e]))

for p, q in zip(rete.W + rete.b, migliore[1]):                       # restore the weights of the best epoch
    p[...] = q
print("\nTEST ACCURACY (once only): %.3f" % rete.accuratezza(X_te, y_te))

previsioni = rete.avanti(X_te).argmax(axis=1)
errori = [(int(a), int(b)) for a, b in zip(y_te, previsioni) if a != b]
print("test errors: %d out of %d;  (true, predicted) pairs:" % (len(errori), len(y_te)), sorted(errori))

epochs run: 42   best epoch: 31
epoch   train loss   val. loss   train acc.   val. acc.
    0       1.2591      1.3788        0.632       0.619
    5       0.1777      0.3299        0.962       0.936
   10       0.0782      0.2590        0.988       0.956
   31       0.0118      0.2334        1.000       0.958
   41       0.0067      0.2348        1.000       0.964

TEST ACCURACY (once only): 0.958
test errors: 15 out of 360;  (true, predicted) pairs: [(1, 8), (2, 1), (3, 8), (3, 9), (4, 1), (4, 8), (6, 1), (8, 1), (8, 1), (8, 5), (8, 7), (8, 9), (8, 9), (9, 3), (9, 5)]

What the code does. Training stops after 42 epochs (10 after the best epoch, the 31st). In the table you can see the typical overfitting gap: the training loss keeps falling (0.012 at the best epoch, 0.007 at the end, accuracy 1.000) while the validation loss stops improving (0.233) and then rises slightly (0.235). The training set has reached 100%: the network has learned to memorize those 1077 examples; validation sits around 96%.

The number that matters, the one on the test set, is 0.958. That is 15 errors out of 360, and many involve morphologically similar digits (an 8 mistaken for a 1 or a 9, for example): a reasonable kind of error for 8×8 pixel images. Restoring the best weights with p[...] = q is necessary because, at the moment of stopping, the current weights are those of 10 epochs after the minimum.

Why should the test set be opened only once? Because every time you look at the test and then change something (the learning rate, the size of the network) you are using it to choose, and the number stops being an honest estimate. Choices are made on the validation set.

The effects of batch size

print("batch size   steps per epoch   val. loss after 20 epochs   val. acc.")
for batch in (8, 32, 128, len(y_tr)):
    rete = Rete([64, 32, 10])
    storia, _ = addestra(rete, epoche=20, batch=batch, lr=0.1, ottimizzatore="sgd")
    print("%10d   %15d   %25.4f   %9.3f" % (batch, int(np.ceil(len(y_tr) / batch)), storia[-1][1], storia[-1][3]))

batch size   steps per epoch   val. loss after 20 epochs   val. acc.
         8               135                     0.2538       0.964
        32                34                     0.2189       0.961
       128                 9                     0.2778       0.947
      1077                 1                     1.1155       0.703

What the code does. With the same number of epochs and the same learning rate, the number of updates changes: 135 steps per epoch with batch 8, 34 with 32, 9 with 128 and just one with the full batch. With the same number of epochs, the full batch made 20 updates in total and has not arrived yet (loss 1.11, accuracy 0.70); batches of 8 and 32 are close to the optimum (0.254 and 0.219), batch 128 is a little behind (0.278). In general, small batches take more and noisier steps (sometimes with a regularizing effect), large batches take precise but fewer steps and make better use of GPUs. Changing the batch size usually requires retuning the learning rate.

The zero-initialization trap

zeri = Rete([64, 32, 10], init="zero")
storia, _ = addestra(zeri, epoche=5, batch=32, lr=0.1, ottimizzatore="sgd")
righe_identiche = np.allclose(zeri.W[0], zeri.W[0][:, [0]])               # are all columns of W1 equal to the first?
print("validation accuracy after 5 epochs: %.3f  (chance with 10 classes would be 0.100)" % storia[-1][3])
print("all 32 hidden neurons have the same input weights:", righe_identiche)

validation accuracy after 5 epochs: 0.100  (chance with 10 classes would be 0.100)
all 32 hidden neurons have the same input weights: True

What the code does. It initializes all the weights to zero and trains. The result is chance-level accuracy (0.1): the network learns nothing useful. The reason is symmetry: if all the neurons of a layer start with the same weights, they receive the same gradients (equation (3) of backpropagation gives the same column for all of them) and stay identical forever. A network with 32 identical hidden neurons is a network with a single neuron. This is why weights are initialized at random, with the He scale for ReLU (layered networks lesson); the biases can start at zero because their gradient depends on the weights.

The same thing with a library

A library automates the three ingredients you wrote by hand: the computation of gradients, the optimizers and the handling of mini-batches. The loop remains, identical in structure. Here we use scikit-learn's MLPClassifier, which is available on the working machine, with the same hyperparameters as our Adam experiment: 32 hidden neurons, learning rate 0.002, mini-batches of 32, at most 200 epochs, stopping after 10 epochs without improvement. Two differences must be stated: scikit-learn's early stopping is based on an internal validation split of 10% of the training set (not on our validation set), and the weight initialization differs from ours.

from sklearn.neural_network import MLPClassifier

mlp = MLPClassifier(hidden_layer_sizes=(32,), solver="adam", learning_rate_init=0.002, batch_size=32,
                    max_iter=200, early_stopping=True, n_iter_no_change=10, random_state=0)
mlp.fit(X_tr, y_tr)
n_par = sum(c.size for c in mlp.coefs_) + sum(i.size for i in mlp.intercepts_)
print("scikit-learn: epochs run %d   internal validation score (10%% of training) %.3f" % (mlp.n_iter_, mlp.best_validation_score_))
print("validation accuracy %.3f   TEST ACCURACY %.3f   parameters: %d" % (mlp.score(X_va, y_va), mlp.score(X_te, y_te), n_par))

scikit-learn: epochs run 26   internal validation score (10% of training) 0.944
validation accuracy 0.967   TEST ACCURACY 0.967   parameters: 2410

What the code does. fit contains the entire loop you wrote: reshuffling at every epoch, mini-batches, forward pass, backward pass and Adam update. The parameters are 2410, as the sum 32·65 + 10·33 = 2080 + 330 that follows from the formula in the layered networks lesson. The test accuracy is 0.967 against our 0.958: the difference is 3 examples out of 360 and is not significant, and it comes from the different initialization, the different validation share and the random sequence of batches. It is a good example of a general rule: two correct implementations of the same model do not give the same number, but close numbers, and small differences must be compared with the variability due to the seed (exercise 3).

The same thing in PyTorch

PyTorch automates the same three ingredients: the graph and backpropagation (backward), the optimizers (torch.optim) and the handling of mini-batches (DataLoader). The loop has the same structure as yours. The listing reuses the data X_tr, y_tr, X_va, y_va, X_te, y_te already prepared above (with np), and was run with PyTorch 2.12.0 on CPU (CUDA_VISIBLE_DEVICES="", 2 threads), twice, with torch.manual_seed(0) and identical output; each run takes from a few seconds to half a minute, depending on the load of the machine.

import torch

def a_tensore(a, dtype=torch.float32):
    return torch.tensor(a, dtype=dtype)

dataset = torch.utils.data.TensorDataset(a_tensore(X_tr), a_tensore(y_tr, torch.long))     # class labels are "long" integers
torch.manual_seed(0)
loader = torch.utils.data.DataLoader(dataset, batch_size=32, shuffle=True)                  # reshuffles every epoch

modello = torch.nn.Sequential(torch.nn.Linear(64, 32), torch.nn.ReLU(), torch.nn.Linear(32, 10))
ottimizzatore = torch.optim.Adam(modello.parameters(), lr=0.002)
funzione_perdita = torch.nn.CrossEntropyLoss()                                              # wants logits, applies softmax itself

def valuta(Xs, ys):
    modello.eval()                                                                          # evaluation mode
    with torch.no_grad():                                                                   # no graph: faster, less memory
        uscita = modello(a_tensore(Xs))
        return funzione_perdita(uscita, a_tensore(ys, torch.long)).item(), float((uscita.argmax(dim=1).numpy() == ys).mean())

migliore = (np.inf, None, 0)
for epoca in range(200):
    modello.train()                                                                         # training mode
    for xb, yb in loader:
        ottimizzatore.zero_grad()                                                           # reset the accumulated gradients
        perdita = funzione_perdita(modello(xb), yb)                                         # forward + loss
        perdita.backward()                                                                  # backward
        ottimizzatore.step()                                                                # update
    perdita_val, acc_val = valuta(X_va, y_va)
    if perdita_val < migliore[0] - 1e-4:
        migliore = (perdita_val, {k: v.clone() for k, v in modello.state_dict().items()}, epoca)
    elif epoca - migliore[2] >= 10:
        break

modello.load_state_dict(migliore[1])
print("epochs run: %d   best epoch: %d   validation loss: %.4f" % (epoca + 1, migliore[2], migliore[0]))
print("PyTorch: validation accuracy %.3f   TEST ACCURACY %.3f   parameters: %d" % (
    valuta(X_va, y_va)[1], valuta(X_te, y_te)[1], sum(p.numel() for p in modello.parameters())))

epochs run: 33   best epoch: 22   validation loss: 0.1431
PyTorch: validation accuracy 0.958   TEST ACCURACY 0.967   parameters: 2410

What the code does, line by line. TensorDataset pairs the features and the labels, and DataLoader produces mini-batches reshuffled at every epoch: it replaces your rng.permutation and your loops over indices. nn.Sequential is the 64 → 32 → 10 network. CrossEntropyLoss receives the logits and applies softmax and cross-entropy internally, which is why the model has no final softmax and the labels must be integers (torch.long). The four lines in the inner loop are the ones you know: zero_grad(), the loss computation, backward() (your indietro) and step() (your passo).

modello.eval() and torch.no_grad() in valuta are two distinct things: the first changes the behavior of certain layers (dropout, batch normalization) between training and evaluation, the second disables the construction of the gradient graph. To save the best weights for early stopping you need a copy (clone()), otherwise the dictionary would point to the same tensors that keep changing.

Training stops after 33 epochs with the best one at index 22 (epochs are counted from 0), validation loss 0.143, validation accuracy 0.958 and test accuracy 0.967, with 2410 parameters (32·65 + 10·33 = 2080 + 330). Our by-hand loop had stopped after 42 epochs, with the best at 31 and 0.958 on the test; scikit-learn gives 0.967. The difference is 3 examples out of 360 and is not significant: it comes from the initialization (PyTorch draws the weights from a uniform distribution of amplitude 1/√inputs, we use the He scale with a normal) and from the random sequence of batches. Three correct implementations of the same model do not give the same number, but close numbers.

Diagnosing a training run that does not work: three questions in order

Sooner or later a network will not learn and you will have to understand why. The loop you have built gives you the signals to do it; it is best to ask three questions, in order, because each one rules out an entire family of causes.

1. Does the training loss go down? If not, the problem is in the optimization or in the code, not in the data. Check in this order: the initial loss (with 10 classes the reference is ln 10 ≈ 2.30; our network gave 2.79 with He initialization, which is not nothing; a value ten times larger or nan means a bug), the learning rate (the test above: too small does not move, too large explodes), the shapes of the tensors, the gradient check of the previous lesson and, in PyTorch, the presence of zero_grad(). A good sanity test is to ask the network to overfit a very small number of examples (say 20): if it cannot memorize them, the code has an error.

2. Does the training loss go down but the validation loss does not (or rises again)? It is overfitting: the network memorizes. The levers are those of the regularization lesson (L2 penalty, early stopping), a smaller network, more data or, for images, data augmentation and dropout techniques, which you will see shortly. If the gap is huge from the very first epochs, also check for leakage: an unrealistically high validation accuracy is more often an error in data preparation than a miracle of the model.

3. Are both low and stuck? It is underfitting: the model does not have enough capacity, or there is not enough information in the features, or the optimization has stalled on a plateau. Try a wider or deeper network, a different learning rate, more epochs, but first ask yourself whether the data really contain the signal you are looking for. No architecture extracts from a feature what is not there.

A habit worth gold: always keep two numbers at hand, the trivial baseline (here, predicting the most frequent class: about 0.10 accuracy) and the expected initial loss (here, of the order of ln 10 ≈ 2.3; measured 2.79). With these two references you recognize at a glance a training run that is not learning or a result too good to be true.

Three exercises to fix the ideas

Exercise 1. Halve the number of hidden neurons (from 32 to 16) and double it (64). Measure, with the same procedure, validation loss and accuracy for the same number of epochs. What do you expect for the size of the gap between training and validation? (More neurons, more capacity: the gap tends to grow, and early stopping intervenes sooner.)

Exercise 2. Add an L2 penalty on the weights to the loss function (regularization lesson) by modifying indietro: to each gradient of W you add (λ/m)·W (not to the biases). Try λ in 0.001, 0.01, 0.1 and observe the effect on the gap between training and validation.

Exercise 3. Repeat the training with five different seeds (change both the network's seed and the one of addestra) and report the mean and standard deviation of the validation accuracy. How large is the variability due to chance? It will serve as a yardstick to decide whether a difference between two configurations is real or noise.

Trade-offs and real-world pitfalls

• Using the test set to choose: hyperparameters, architecture, number of epochs. They are chosen on the validation set; the test is opened once, at the end.
• Standardizing before the split: leakage. The scaler learns from the training set alone.
• Forgetting zero_grad() in PyTorch: the gradients accumulate and training becomes unstable.
• Forgetting model.eval() and torch.no_grad() in evaluation: with dropout or batch norm the results are wrong, and memory grows needlessly.
• Not shuffling the data: if the dataset is sorted by class, every mini-batch contains a single digit and training is a disaster.
• Labels of the wrong type: CrossEntropyLoss wants integers (torch.long), not floats; and logits, not probabilities.
• Monitoring only the training loss: it always goes down. Overfitting shows in the gap with validation.
• Interpreting a single run: repeat with different seeds before saying that an idea "works"; differences of a couple of examples on the test are noise.
• Learning rate not matched to the batch or the optimizer: SGD and Adam live on different learning-rate scales (0.1 versus 0.002 here).
• Restoring the best weights: if you forget, the final model is the one after overfitting.

The ten entries are divided by the damage they do. The first two (test used to choose, standardization before the split) produce an optimistic number without any visible error. The entries on zero_grad, eval, no_grad and on the label type concern PyTorch: they describe the documented behavior of the library, and the PyTorch loop above ran correctly, but I did not provoke each of these mistakes one by one to show its effect. The others are habits of method (shuffling, looking at the gap, repeating with several seeds, restoring the best weights) that cost a few minutes and protect you from wrong conclusions.

Frequently asked questions

How many epochs are needed? It is not decided in advance: you set a high number and use early stopping on the validation set. In this example about thirty were enough.

Why does validation not coincide with the test? Because validation has already influenced your choices (when to stop, which learning rate). The test must remain uncontaminated.

Is it a good result? 0.96 on 10 classes with 1077 training examples and a tiny network is a good result for a baseline, but a convolutional network, which you will see shortly, exploits the structure of images and typically does better. There is no absolute "good": you compare with a baseline (here, the most frequent class would give about 0.10) and with alternative models.

Should I always use Adam? It is a good starting choice because it requires less tuning, but SGD with momentum and a good learning-rate schedule is still widely used and, in some vision problems, generalizes better. It depends on the problem.

What you can do now

You can prepare data without leakage, write a mini-batch training loop with three optimizers, read learning curves and stop training at the right moment, opening the test set just once. You also know what happens with a wrong learning rate, with batches that are too large or with weights all at zero, and you have seen your loop reproduced by a library with close results. This lesson builds on backpropagation and layered networks; the next step is convolution, which applies these ideas to data with structure, images.

Other lessons in the series that link to this one: Entropy, Cross-Entropy and KL Divergence: Why Loss Functions Are All the Same Thing, Training a LeNet-5 style CNN with PyTorch.