Your first neural network in PyTorch

Lesson 04 built the training loop by hand in numpy: a forward pass, a loss, gradients, a step against the gradient. This lesson does the same thing in PyTorch, where the gradients are computed for you, and uses it to train a small neural network that reads handwritten digits. The whole script runs in under five seconds on a laptop CPU. It is for anyone who has used scikit-learn’s fit and predict and wants to see what is inside the box when the model is a network.

The dataset is sklearn.datasets.load_digits: 1,797 images of 8 by 8 pixels, each a digit from 0 to 9, each pixel a grey level from 0 to 16. It ships with scikit-learn so nothing is downloaded. It is small enough that you can print an image as a grid of numbers and recognise the digit, which is a useful sanity check that bigger datasets do not allow.

Tensors and automatic gradients

A tensor is PyTorch’s array. It behaves like a numpy array with two additions: it can live on an accelerator (a GPU, or Apple’s MPS backend), and it can record the operations applied to it so that gradients can be computed afterwards. That second feature, autograd, is what makes the hand-written gradient code from lesson 04 unnecessary.

import torch

torch.manual_seed(0)
a = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])   # shape (2, 3)
w = torch.randn(3, 4)                                   # shape (3, 4)
b = torch.zeros(4)
out = a @ w + b                                         # matrix product, shape (2, 4)
print("a.shape =", tuple(a.shape), " w.shape =", tuple(w.shape), " out.shape =", tuple(out.shape))
print("dtype:", out.dtype, " device:", out.device)

x = torch.tensor(2.0, requires_grad=True)
y = x ** 2 + 3 * x
y.backward()
print("dy/dx at x=2:", x.grad.item())   # 2x + 3 = 7

a.shape = (2, 3)  w.shape = (3, 4)  out.shape = (2, 4)
dtype: torch.float32  device: cpu
dy/dx at x=2: 7.0

The first half is the shape arithmetic you will do constantly: a batch of 2 rows with 3 features, multiplied by a 3 by 4 weight matrix, gives 2 rows with 4 outputs. The bias b has shape (4,) and is added to every row by broadcasting. The tensor shapes reference covers the rules in full.

The second half is autograd. Setting requires_grad=True tells PyTorch to record every operation involving x. Calling backward() on the scalar y walks that record in reverse and fills x.grad with dy/dx. For a network, x is replaced by thousands of weights and y by the loss, and the same call fills in every gradient at once.

Loading the digits as tensors

The data preparation is ordinary scikit-learn: a stratified split, then standardisation using statistics from the training set only (the rule from lesson 02). The split here is called validation rather than test because you will watch it during training to decide when to stop and which settings to keep, which means it is being used for decisions. A true test set would be held back further.

import numpy as np
import torch
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split

digits = load_digits()
X, y = digits.data.astype(np.float32), digits.target
print("images:", X.shape, " pixel range:", X.min(), "to", X.max())

X_train, X_val, y_train, y_val = train_test_split(
    X, y, test_size=0.2, stratify=y, random_state=0
)
mean, std = X_train.mean(), X_train.std()
X_train = (X_train - mean) / std      # scale with training statistics only
X_val = (X_val - mean) / std

X_train_t = torch.from_numpy(X_train)
y_train_t = torch.from_numpy(y_train)
X_val_t = torch.from_numpy(X_val)
y_val_t = torch.from_numpy(y_val)
print("train:", tuple(X_train_t.shape), "val:", tuple(X_val_t.shape), "labels dtype:", y_train_t.dtype)

images: (1797, 64)  pixel range: 0.0 to 16.0
train: (1437, 64) val: (360, 64) labels dtype: torch.int64

Two dtype details will save you an error message later. Features are float32, which is PyTorch’s default and what nn.Linear expects (numpy’s default is float64, hence the astype). Labels for classification are integers, int64, because the loss function uses them as indices into the output columns. The 64 pixels of each image are flattened into one row, so the network sees a vector, not a grid. A convolutional network, which keeps the grid, is the subject of the MNIST CNN project.

The model as an nn.Module

A multilayer perceptron (MLP) is a stack of linear layers with a non-linearity between them. With one hidden layer it computes h = relu(x @ W1 + b1) and then scores = h @ W2 + b2. The diagram shows the shapes for this lesson’s network.

The multilayer perceptron used in this lesson. Each box is a tensor operation, each wire carries a tensor of the shape shown, for a batch of 32 images.

Without the ReLU the two linear layers would collapse into one (a product of matrices is a matrix), and the model would be the logistic regression of lesson 03 with extra steps. ReLU, max(0, h), is the simplest function that breaks that collapse. It zeroes negative activations and passes positive ones unchanged, and its gradient is 0 or 1, which keeps training stable.

The outputs are called logits: one raw, unbounded score per class. They are not probabilities. Cross-entropy loss applies softmax internally to turn them into probabilities and then takes the negative log of the probability assigned to the true class. Keeping softmax inside the loss rather than inside the model is both numerically safer and the convention every PyTorch tutorial follows, so do not add a softmax layer at the end of a classifier.

In PyTorch a model is a class that inherits from nn.Module, builds its layers in __init__, and defines the forward pass in forward. The module tracks every parameter inside the layers you assign to self, so the optimiser can find them with model.parameters().

The training loop

The block below is the whole thing: model, loss, optimiser and loop. Read the loop body against the five steps from lesson 04. Forward (logits = model(...)), loss, zero the old gradients, backward, step. The only new element is minibatching: each epoch shuffles the training rows and walks through them 32 at a time, so there are 45 updates per epoch rather than one.

import numpy as np
import torch
from torch import nn
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split

torch.manual_seed(0)
np.random.seed(0)

digits = load_digits()
X, y = digits.data.astype(np.float32), digits.target
X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2, stratify=y, random_state=0)
mean, std = X_train.mean(), X_train.std()
X_train_t = torch.from_numpy((X_train - mean) / std)
X_val_t = torch.from_numpy((X_val - mean) / std)
y_train_t = torch.from_numpy(y_train)
y_val_t = torch.from_numpy(y_val)


class Mlp(nn.Module):
    def __init__(self, n_in: int = 64, n_hidden: int = 64, n_out: int = 10):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(n_in, n_hidden),
            nn.ReLU(),
            nn.Linear(n_hidden, n_out),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)          # raw scores (logits), one per class


model = Mlp()
n_params = sum(p.numel() for p in model.parameters())
print("parameters:", n_params)

loss_fn = nn.CrossEntropyLoss()
optimiser = torch.optim.SGD(model.parameters(), lr=0.05)
batch_size = 32
n_train = len(X_train_t)

for epoch in range(1, 21):
    model.train()
    perm = torch.randperm(n_train)
    running = 0.0
    for start in range(0, n_train, batch_size):
        idx = perm[start:start + batch_size]
        logits = model(X_train_t[idx])
        loss = loss_fn(logits, y_train_t[idx])
        optimiser.zero_grad()
        loss.backward()
        optimiser.step()
        running += loss.item() * len(idx)
    train_loss = running / n_train

    model.eval()
    with torch.no_grad():
        val_logits = model(X_val_t)
        val_loss = loss_fn(val_logits, y_val_t).item()
        val_acc = (val_logits.argmax(dim=1) == y_val_t).float().mean().item()
    print(f"epoch {epoch:2d}  train loss {train_loss:.3f}  val loss {val_loss:.3f}  val acc {val_acc:.3f}")

parameters: 4810
epoch  1  train loss 1.895  val loss 1.382  val acc 0.850
epoch  2  train loss 0.966  val loss 0.685  val acc 0.906
epoch  3  train loss 0.519  val loss 0.458  val acc 0.914
epoch  4  train loss 0.357  val loss 0.345  val acc 0.936
epoch  5  train loss 0.276  val loss 0.287  val acc 0.950
epoch  6  train loss 0.231  val loss 0.252  val acc 0.950
epoch  7  train loss 0.198  val loss 0.236  val acc 0.956
epoch  8  train loss 0.175  val loss 0.217  val acc 0.953
epoch  9  train loss 0.161  val loss 0.198  val acc 0.956
epoch 10  train loss 0.145  val loss 0.180  val acc 0.961
epoch 11  train loss 0.135  val loss 0.176  val acc 0.958
epoch 12  train loss 0.123  val loss 0.163  val acc 0.969
epoch 13  train loss 0.117  val loss 0.167  val acc 0.961
epoch 14  train loss 0.109  val loss 0.161  val acc 0.961
epoch 15  train loss 0.104  val loss 0.154  val acc 0.964
epoch 16  train loss 0.098  val loss 0.148  val acc 0.961
epoch 17  train loss 0.093  val loss 0.141  val acc 0.972
epoch 18  train loss 0.089  val loss 0.139  val acc 0.967
epoch 19  train loss 0.085  val loss 0.144  val acc 0.961
epoch 20  train loss 0.081  val loss 0.136  val acc 0.969

The parameter count is 64 x 64 + 64 for the first layer and 64 x 10 + 10 for the second, 4,810 in all. Each epoch is 45 minibatches and the whole run takes about two seconds on CPU. Exact losses will differ slightly between machines and PyTorch versions even with the seed set, because floating-point reductions are not bit-identical across hardware, but the shape of the curve will be the same.

Cross-entropy loss per epoch on the training and validation splits, from the printed numbers above.

Reading the curves

Both losses fall steeply for five epochs and then flatten. The validation loss sits above the training loss from about epoch 4 onwards and the gap widens slowly: the model fits its training rows a little better than it fits unseen rows, which is normal and small here. If the validation loss had started rising while the training loss kept falling, that would be overfitting, and the response would be to stop earlier, shrink the hidden layer, add weight decay or get more data. Twenty epochs of this network on this data does not reach that point. Validation accuracy ends near 97 %. Be honest about what that shows: a scikit-learn logistic regression on the same standardised split also scores about 97 %, because these digits are small, clean and close to linearly separable. The hidden layer does not buy accuracy here. What this lesson buys is the machinery, which is unchanged when the data is images of 224 by 224 pixels and the model has a hundred layers, where a linear model is hopeless.

The validation loss in epoch 1 is lower than the training loss, which surprises people. The training loss is averaged over the epoch while the weights are still moving, so it includes the bad early batches. The validation loss is measured once, at the end of the epoch, with the improved weights.

Three lines in the loop are easy to forget and each produces a quiet bug rather than an error. optimiser.zero_grad() is needed because PyTorch accumulates gradients into .grad rather than replacing them, so without it every step uses the sum of all gradients so far. model.eval() and model.train() switch the behaviour of layers such as dropout and batch normalisation. This network has neither, so they do nothing here, but the habit matters the first time you add one. torch.no_grad() tells autograd not to record the validation forward pass, which saves memory and time. The PyTorch training loop reference lists these and the other standard lines in one place.

What changes at scale

Everything in this lesson scales up without changing shape. A larger network has more layers in nn.Sequential. A larger dataset is wrapped in a DataLoader that produces the minibatches, instead of the manual randperm slicing used here. A GPU is used by calling .to(device) on the model and on each batch. torch.optim.Adam replaces SGD in most modern work because it adapts the step size per parameter, at the cost of a little memory. The loop body is otherwise the same five lines, and the same curves are what you watch. The deep learning hub and the PyTorch platform page take it from here.

Exercise

Change n_hidden to 8, then to 256, and run 20 epochs each time. Record the final train and validation loss for all three sizes. Which one shows the largest gap between the two, and does the smallest network underfit.
Replace torch.optim.SGD(..., lr=0.05) with torch.optim.Adam(model.parameters(), lr=1e-3) and compare the curve after 5 epochs. Then try lr=0.5 with SGD and describe what happens to the loss.
After training, find the validation images the model gets wrong: wrong = (val_logits.argmax(dim=1) != y_val_t).nonzero(). Print one of them as an 8 by 8 grid with X_val[i].reshape(8, 8).round(1) and decide whether you would have got it right.

Summary

A tensor is an array that can live on an accelerator and can record operations so that backward() computes gradients automatically.
An nn.Module declares its layers in __init__ and its computation in forward. Parameters inside the layers are tracked for you.
A multilayer perceptron is linear, ReLU, linear. The non-linearity is what makes it more than a single linear model. Outputs are logits, and cross-entropy loss applies the softmax itself.
The loop is forward, loss, zero_grad, backward, step, run over shuffled minibatches, with a no_grad validation pass at the end of each epoch.
Watch train and validation loss together. A widening gap with rising validation loss is overfitting. A loss stuck at ln(number of classes) means nothing is being learned.