An MNIST-class CNN in PyTorch on CPU

What you will build

This project is the convolutional network everyone writes once: two convolution layers, two pooling layers, two fully connected layers, ten output classes. The dataset is scikit-learn’s bundled load_digits, 1,797 handwritten digits at 8 by 8 pixels, upsampled to 16 by 16 so the convolutions have something to do. No download, no torchvision, and the whole run takes about a second of training on a laptop CPU. MNIST itself is 70,000 images at 28 by 28 and needs a download, so it is left as the first thing to try afterwards. Nothing in the model changes except one number.

You will write the data loader, the model, the training loop and the evaluation by hand, print loss and test accuracy per epoch, save the weights, reload them in a second program and print a confusion matrix. The confusion matrix is the interesting bit. A single accuracy number says the model is 95 % right. The matrix says which digit it keeps mistaking for which, and that tells you what to do next.

This follows on from Your first neural network in PyTorch and the training loop lesson. The PyTorch training loop cheat sheet has the loop skeleton on one page.

What you need

• Python 3.10 or later with torch (CPU build is enough), numpy and scikit-learn. We tested with PyTorch 2.14.1 and scikit-learn 1.9.1.
• No GPU. The code pins device = torch.device("cpu") so that two runs print the same numbers. The Things to try section shows the one-line change for Apple Silicon or CUDA.
• About 40 MB of memory for the data and model. The saved weights are 155 kB.

Both Python blocks were run from an empty directory with the Field Guide’s pipeline environment on an Apple Silicon Mac mini. Timings are given with each block.

Design

Why upsample 8 by 8 to 16 by 16? Two rounds of 2 by 2 max pooling divide each spatial dimension by four. From 8 by 8 that leaves a 2 by 2 map, which barely counts as an image. From 16 by 16 it leaves 4 by 4, enough for the second convolution to see strokes rather than single pixels. Bilinear interpolation adds no information, it only gives the kernels room. The same code with size=28 and MNIST inputs needs only the fc1 input size changed to 32 * 7 * 7.

Why these layers? conv1 learns 16 filters of 3 by 3 pixels: edges and stroke ends. conv2 learns 32 combinations of those over a wider field. Padding of 1 keeps the spatial size through each convolution so that only pooling shrinks it. Each ReLU is followed by pooling, which keeps the strongest response in each 2 by 2 window and halves the work for the next layer. The classifier is one hidden layer of 64 units. The whole thing has 38,282 parameters, printed by the training block.

Why a plain loop and not DataLoader? The dataset is 1,437 training images that fit in one tensor. Shuffling with torch.randperm and slicing it in batches of 32 is simpler to read and faster than a DataLoader at this size. Switch to DataLoader when the data no longer fits in memory or when you need transforms per sample.

Why stratify the split? train_test_split(..., stratify=y) keeps the same proportion of each digit in train and test. With 360 test images and ten classes, an unlucky random split could leave one digit with 25 test examples and another with 45, which makes per-class accuracy noisy.

Code

Train

Save this as train.py. The loop reports the mean training loss over the epoch and the test accuracy after it. Evaluation runs under torch.no_grad() with the model in eval() mode, which changes nothing for this model (no dropout or batch norm) but is the habit to keep.

# train.py: a two-layer CNN on 16x16 digit images, trained on CPU.
import time
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

torch.manual_seed(0)
np.random.seed(0)
device = torch.device("cpu")   # deterministic and fast enough at this size


def load_images(size: int = 16):
    """1,797 greyscale digits, 8x8 with values 0-16, upsampled to size x size in [0, 1]."""
    digits = load_digits()
    x = torch.tensor(digits.images, dtype=torch.float32).unsqueeze(1) / 16.0  # (N, 1, 8, 8)
    x = F.interpolate(x, size=(size, size), mode="bilinear", align_corners=False)
    y = torch.tensor(digits.target, dtype=torch.long)
    return x, y


class SmallCNN(nn.Module):
    def __init__(self, n_classes: int = 10):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 16, kernel_size=3, padding=1)   # 16 x 16 x 16
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)  # 32 x 8 x 8 after pool
        self.fc1 = nn.Linear(32 * 4 * 4, 64)
        self.fc2 = nn.Linear(64, n_classes)

    def forward(self, x):
        x = F.max_pool2d(F.relu(self.conv1(x)), 2)   # 16 x 8 x 8
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)   # 32 x 4 x 4
        x = torch.flatten(x, 1)
        x = F.relu(self.fc1(x))
        return self.fc2(x)


x, y = load_images()
idx_train, idx_test = train_test_split(
    np.arange(len(y)), test_size=0.2, random_state=0, stratify=y.numpy()
)
x_train, y_train = x[idx_train].to(device), y[idx_train].to(device)
x_test, y_test = x[idx_test].to(device), y[idx_test].to(device)

model = SmallCNN().to(device)
n_params = sum(p.numel() for p in model.parameters())
optimiser = torch.optim.Adam(model.parameters(), lr=3e-3)
batch_size = 32
epochs = 5
print(f"train {len(y_train)}, test {len(y_test)}, image {tuple(x.shape[1:])}, params {n_params:,}")

t0 = time.perf_counter()
for epoch in range(1, epochs + 1):
    model.train()
    perm = torch.randperm(len(y_train))
    total_loss = 0.0
    for start in range(0, len(perm), batch_size):
        batch = perm[start:start + batch_size]
        logits = model(x_train[batch])
        loss = F.cross_entropy(logits, y_train[batch])
        optimiser.zero_grad()
        loss.backward()
        optimiser.step()
        total_loss += loss.item() * len(batch)
    model.eval()
    with torch.no_grad():
        test_acc = (model(x_test).argmax(1) == y_test).float().mean().item()
    print(f"epoch {epoch}  train loss {total_loss / len(y_train):.4f}  test acc {test_acc:.4f}")

print(f"training time {time.perf_counter() - t0:.1f} s on {device}")
torch.save(model.state_dict(), "digits_cnn.pt")
print("saved digits_cnn.pt")

Expected output. The whole script took 3.3 s on the Mac mini, of which 0.8 s was training. The rest is importing PyTorch.

train 1437, test 360, image (1, 16, 16), params 38,282
epoch 1  train loss 1.7641  test acc 0.8194
epoch 2  train loss 0.3834  test acc 0.9111
epoch 3  train loss 0.2169  test acc 0.9167
epoch 4  train loss 0.1518  test acc 0.9306
epoch 5  train loss 0.0955  test acc 0.9528
training time 0.8 s on cpu
saved digits_cnn.pt

Two runs on the same machine printed identical numbers, because every source of randomness (the split, the weight initialisation, the batch order) is seeded and the computation is on CPU. On MPS or CUDA expect the numbers to drift in the third decimal place.

The first epoch’s mean loss of 1.76 is high because it averages over the whole epoch, including the first batches where the network is guessing (a uniform guess over ten classes gives a loss of ln 10, about 2.30). By the end of epoch 1 the model is already at 82 % on the test set. The loss is still falling at epoch 5 and so is the gap to 100 %, which says more epochs would help. That is deliberate: the Things to try section asks you to find out how many.

Evaluate with a confusion matrix

Save this as evaluate.py. It restates the model class, rebuilds the same test split, loads the saved weights and prints the matrix. If digits_cnn.pt is missing it trains the same model with the same recipe first, so the block runs on its own. In a real project the class would live in a module both scripts import.

# evaluate.py: reload the saved weights and print a confusion matrix.
from pathlib import Path

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.metrics import confusion_matrix
from sklearn.model_selection import train_test_split

torch.manual_seed(0)


class SmallCNN(nn.Module):
    def __init__(self, n_classes: int = 10):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 16, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)
        self.fc1 = nn.Linear(32 * 4 * 4, 64)
        self.fc2 = nn.Linear(64, n_classes)

    def forward(self, x):
        x = F.max_pool2d(F.relu(self.conv1(x)), 2)
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = torch.flatten(x, 1)
        x = F.relu(self.fc1(x))
        return self.fc2(x)


digits = load_digits()
x = torch.tensor(digits.images, dtype=torch.float32).unsqueeze(1) / 16.0
x = F.interpolate(x, size=(16, 16), mode="bilinear", align_corners=False)
y = digits.target
_, idx_test = train_test_split(np.arange(len(y)), test_size=0.2, random_state=0, stratify=y)

model = SmallCNN()
if not Path("digits_cnn.pt").exists():
    # Same recipe as train.py, so this block also works on its own.
    idx_train, _ = train_test_split(np.arange(len(y)), test_size=0.2, random_state=0, stratify=y)
    xt, yt = x[idx_train], torch.tensor(y[idx_train])
    optimiser = torch.optim.Adam(model.parameters(), lr=3e-3)
    for epoch in range(5):
        perm = torch.randperm(len(yt))
        for start in range(0, len(perm), 32):
            b = perm[start:start + 32]
            loss = F.cross_entropy(model(xt[b]), yt[b])
            optimiser.zero_grad()
            loss.backward()
            optimiser.step()
    torch.save(model.state_dict(), "digits_cnn.pt")
model.load_state_dict(torch.load("digits_cnn.pt", map_location="cpu"))
model.eval()
with torch.no_grad():
    pred = model(x[idx_test]).argmax(1).numpy()
truth = y[idx_test]

cm = confusion_matrix(truth, pred, labels=range(10))
print("rows = true digit, columns = predicted digit")
print("     " + " ".join(f"{c:3d}" for c in range(10)))
for r in range(10):
    print(f"{r:3d}: " + " ".join(f"{v:3d}" for v in cm[r]))
per_class = cm.diagonal() / cm.sum(1)
print("per-class accuracy:", " ".join(f"{a:.2f}" for a in per_class))
print(f"overall accuracy: {(pred == truth).mean():.4f} on {len(truth)} test images")

worst = [(int(cm[r, c]), r, c) for r in range(10) for c in range(10) if r != c]
worst.sort(reverse=True)
print("most common confusions (count, true, predicted):", worst[:3])

Expected output, 1.7 s on the Mac mini with the saved weights present (3.2 s when it has to train first):

rows = true digit, columns = predicted digit
       0   1   2   3   4   5   6   7   8   9
  0:  35   0   0   0   1   0   0   0   0   0
  1:   0  36   0   0   0   0   0   0   0   0
  2:   0   1  34   0   0   0   0   0   0   0
  3:   0   0   0  36   0   0   0   1   0   0
  4:   0   1   0   0  34   0   0   0   0   1
  5:   0   0   0   1   0  36   0   0   0   0
  6:   0   1   0   0   0   0  35   0   0   0
  7:   0   0   0   0   0   0   0  36   0   0
  8:   0   7   0   0   0   1   0   1  26   0
  9:   0   1   0   0   0   0   0   0   0  35
per-class accuracy: 0.97 1.00 0.97 0.97 0.94 0.97 0.97 1.00 0.74 0.97
overall accuracy: 0.9528 on 360 test images
most common confusions (count, true, predicted): [(7, 8, 1), (1, 9, 1), (1, 8, 7)]

Read the matrix along the rows. Nine of the ten digits are at 94 % or better. The digit 8 is at 74 %, and seven of its 35 test images were called 1. That one cell accounts for almost half of all the errors. At 8 by 8 pixels a narrow 8 and a 1 share most of their ink, and bilinear upsampling does not add the loops back. The column for 1 also collects a stray 2, 4, 6 and 9, so the model has a general bias towards predicting 1 on thin shapes. More epochs fix most of this (try it), and so does a wider first layer. The matrix told you where to look, which is something the accuracy number alone never does.

Run

pip install torch numpy scikit-learn
python train.py
python evaluate.py

Both scripts write to the current directory, so run them from the same place. train.py overwrites digits_cnn.pt on each run.

Test

The two outputs above are the test. The checks that matter:

• The printed parameter count is 38,282. If it differs, a layer size is wrong.
• Loss falls every epoch and test accuracy ends at 0.95 or above. A flat loss of about 2.30 means the optimiser is not stepping (the usual cause is a missing optimiser.step() or zero_grad() in the wrong place).
• The confusion matrix rows each sum to 36 or 35 (the stratified split gives each digit 35 or 36 test images) and the diagonal dominates.
• evaluate.py gives the same accuracy as the last epoch of train.py. If it does not, the test split or the preprocessing in the two files has diverged.

Things to try

• Train for 20 epochs instead of 5 and watch the 8 versus 1 cell. Then plot training loss against test accuracy and find the epoch after which accuracy stops improving while loss keeps falling. That gap is the beginning of overfitting, covered in Evaluation that does not lie.
• Change device to torch.device("mps" if torch.backends.mps.is_available() else "cpu") on a Mac, or "cuda" on an NVIDIA machine, and time it. At this size the CPU is probably faster, because the per-batch overhead of launching work on the accelerator outweighs the arithmetic. Make the batch size 256 and check again.
• Replace load_digits with MNIST through torchvision.datasets.MNIST (a 10 MB download). Set size=28, change fc1 to nn.Linear(32 * 7 * 7, 64), and expect about 98.5 % after 5 epochs on CPU in a few minutes.
• Add nn.Dropout(0.3) before fc2 and nn.BatchNorm2d after each convolution. Note that both change behaviour between train() and eval() mode, which is why the loop switches modes.
• Save the model with torch.save(model, ...) instead of its state_dict() and see what breaks when you rename the class. Then go back to saving the state dictionary.
• Serve the model from the FastAPI app in the tabular classifier project, accepting a 16 by 16 list of lists. The loading code goes in the lifespan hook, the forward pass in the route.

Summary

• Two convolutions with pooling and a small classifier reach 95 % on 16 by 16 digits in under a second of CPU time.
• Padding keeps spatial size through convolutions. Pooling is where the image shrinks. Track the shapes and the fc1 input size falls out.
• Seed everything and pin the device if you want runs to be reproducible to the last digit.
• Save the state dictionary, restate the class in the loading program, and check that the reloaded model gives the same test accuracy.
• A confusion matrix turns one accuracy number into a list of specific problems. Here it points straight at 8 versus 1.