Recognising handwritten digits

The tasks in the previous lessons were tiny: 50 house prices, 4 XOR points. This lesson handles real data: 1,797 images of handwritten digits, with the network recognising which of 0 to 9 each one shows.

This is a classification task. House price prediction outputs a continuous number, which is called regression; recognising digits means choosing one of 10 categories, so both the form of the output and the loss function change. This lesson covers what they change to and why.

The data is the handwritten digits dataset that comes with scikit-learn; it's there once installed, with nothing to download:

uv add torch scikit-learn
python digit_classifier.py

What the data looks like

digits = load_digits()

数据：1797 张图，每张 8×8 像素，像素值 0～16
第一张图，标签是 0：
        ::**==      
        **##++##::  
      ..##..  ++==  
      ::**    ====  
      ::==    ====  
      ::++    **--  
      ..##::++**    
        --**++      

Each image is 8×8, 64 pixels in total, and each pixel is an integer from 0 to 16, larger meaning darker. The script draws the shading with a few characters, and you can see it's a 0.

The network's input is these 64 numbers, flattened into a row. Dividing by 16 scales them to between 0 and 1, for the same reason as Lesson 2's standardisation:

X = torch.tensor(digits.data / 16.0, dtype=torch.float32)
y = torch.tensor(digits.target, dtype=torch.long)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, random_state=0)

训练集 1347 张，测试集 450 张

A quarter of the images are held out as the test set, never shown to the network during training. At the end, it's used to check whether the network has really learned to recognise digits or has just memorised the images in the training set. It's the same idea as the evaluation set in Module 06: the exam questions mustn't appear in the practice questions.

The network: output 10 scores

model = nn.Sequential(
    nn.Linear(64, 64),  # 64 个像素 → 64 个隐藏神经元
    nn.ReLU(),
    nn.Linear(64, 10),  # → 10 个输出，分别对应数字 0～9
)

模型共 4810 个参数

nn.Sequential chains several layers in order, saving you from writing a class. Count the parameters: the first layer has 64×64 weights plus 64 biases, 4,160; the second has 64×10 plus 10, 650; 4,810 in total.

Two things differ from before.

The activation function is now ReLU. It's even simpler than tanh: output the input unchanged if it's greater than 0, otherwise output 0. It's fast to compute, and trains better than tanh in networks with many layers; it and its variants are the most common in today's networks.

The output is 10 numbers, one per digit. They're called scores (logits), and can be any positive or negative number; whichever is largest is the digit the network thinks it is.

softmax: turning scores into probabilities

10 scores aren't easy to interpret directly. softmax turns them into 10 probabilities: first take the exponential of each score (making them positive), then divide by their sum (making them add up to 1). The higher the original score, the higher the probability.

probs = torch.softmax(model(X_test[:1]), dim=1)[0]

After training, the probabilities for the first image in the test set:

测试集第一张图（标签 2）的预测概率：
  0:0.00  1:0.00  2:1.00  3:0.00  4:0.00  5:0.00  6:0.00  7:0.00  8:0.00  9:0.00

The network is almost 100% sure it's a 2, and it is indeed a 2.

You saw exactly the same thing in Module 01, Lesson 2: what an LLM outputs at each step is a score for every token in its vocabulary, turned into probabilities by softmax, from which the next token is chosen. The only difference is that an LLM's "categories" number over a hundred thousand.

Cross-entropy: the loss for classification

Regression uses mean squared error; classification uses cross-entropy. It's simple to compute: look at the probability the network gives the correct answer, take its logarithm, and negate it.

正确答案的概率    损失 = -log(概率)
      1.00              0.00
      0.90              0.11
      0.50              0.69
      0.10              2.30
      0.01              4.61

The closer the probability of the correct answer is to 1, the closer the loss is to 0; the smaller the probability, the larger the loss, and it rises fast. A network that is "confidently wrong" gets punished heavily.

One number worth remembering: at the start of training the network knows nothing, the probabilities of the 10 categories are all about 0.1, and the loss is about -log(0.1) = 2.30. If you train a 10-class network and the loss at the first step is far from 2.3, something is probably wrong.

PyTorch's nn.CrossEntropyLoss combines softmax and cross-entropy into one computation, so the model's last layer outputs scores directly; don't add a softmax yourself.

loss_fn = nn.CrossEntropyLoss()  # 分类问题用交叉熵

Training GPT in Module 09 uses it too, with every token in the vocabulary as a category.

Mini-batch training

In the previous lessons every step computed the gradient on all the data. This lesson takes 64 images at a time instead:

optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
for epoch in range(1, 31):
    # 每一轮把训练集打乱，每次取 64 张图更新一次参数
    order = torch.randperm(len(X_train))
    for i in range(0, len(X_train), 64):
        idx = order[i:i + 64]
        loss = loss_fn(model(X_train[idx]), y_train[idx])
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

This is called mini-batch training. With 1,347 images and 64 per batch, the parameters are updated 22 times per round. One pass through all the training data is called an epoch.

Why not use all the data at once? With a large dataset it simply won't fit in memory; LLM training data runs to trillions of tokens. And although mini-batch gradients are somewhat noisy, they update more often, which usually learns faster than updating once per epoch. The order is shuffled at the start of every epoch so each batch is a random handful, and the network can't come to rely on the order of the data.

The optimiser has changed to Adam too. Building on SGD, it adjusts the step size for each parameter individually: parameters that keep moving in the same direction get bigger steps, and ones that swing back and forth get smaller steps. It automatically eases much of Lesson 2's "one learning rate can't serve both directions" problem. In practice, Adam and its variant AdamW are the most common default choices.

Training results

训练：
  第  1 轮  最后一批的损失 0.1804  训练集准确率 87.8%  测试集准确率 87.6%
  第  2 轮  最后一批的损失 0.0335  训练集准确率 92.7%  测试集准确率 92.4%
  第  5 轮  最后一批的损失 0.3066  训练集准确率 97.6%  测试集准确率 96.7%
  第 10 轮  最后一批的损失 0.0285  训练集准确率 98.7%  测试集准确率 96.0%
  第 20 轮  最后一批的损失 0.0035  训练集准确率 99.7%  测试集准确率 97.3%
  第 30 轮  最后一批的损失 0.0001  训练集准确率 100.0%  测试集准确率 97.8%

On my computer, the whole script finished in under 5 seconds.

Two things are worth looking at.

First, the "loss of the last batch" jumps around: 0.03 at epoch 2, but 0.31 at epoch 5. That's mini-batch noise: each batch has only 64 images (the last has only 1347 - 21×64 = 3), and if it happens to draw a few hard ones, the loss is high. To judge how well training is going, look at metrics over the whole dataset, not the loss of one batch.

Second, training set accuracy reached 100% and test set accuracy 97.8%. The network gets every image it has seen right, and some it hasn't seen wrong. The next lesson is devoted to that gap.

Where it goes wrong

Accuracy is just one number. To see where the network goes wrong, look at the confusion matrix: rows are the true digits, columns are the network's predictions, the diagonal holds the correct answers and everything else is a mistake.

测试集 450 张里错了 10 张。混淆矩阵（行是真实数字，列是预测数字）：
       0   1   2   3   4   5   6   7   8   9
  0   37   .   .   .   .   .   .   .   .   .
  1    .  42   .   .   .   .   .   .   1   .
  2    .   .  44   .   .   .   .   .   .   .
  3    .   .   1  44   .   .   .   .   .   .
  4    .   .   .   .  38   .   .   .   .   .
  5    .   .   .   .   .  47   .   .   .   1
  6    .   1   .   .   .   .  51   .   .   .
  7    .   1   .   .   .   .   .  47   .   .
  8    .   2   1   .   .   .   .   .  45   .
  9    .   .   .   .   .   1   .   1   .  45

Of the 10 errors, 4 were misread as 1 (one 6, one 7 and two 8s). Draw one of them:

一张判错的图：真实是 7，模型认为是 1
      ..####::      
        ::++==      
          **--..    
        ==######++  
      ++####--      
        ##==        
        ##..        
      ..##          

This 7 has an extra stroke across the middle, and its lower half is a single vertical line. At 8×8 resolution, calling it a 1 isn't outrageous. Looking at the misclassified samples is the best way to understand a model: sometimes it's the model's fault, sometimes the data itself is ambiguous, and sometimes the label is even wrong. It's the same as looking at failed cases in Module 06's evaluation.

Exercises

1. Change the hidden layer from 64 neurons to 16 and to 256. How do training and test set accuracy change?
2. Replace Adam with SGD, also at a learning rate of 0.01. What is the test set accuracy after 30 epochs? What learning rate does SGD need to catch up with Adam?
3. Before training starts, compute the loss over the whole training set once. Is it close to 2.30?

Self-check

1. How does the output of a classification task differ from regression? What does softmax do?

Regression outputs one continuous number. Classification outputs a score for each category; softmax takes the exponential of these scores and divides by their sum, turning them into probabilities that add up to 1, with higher-scoring categories getting higher probabilities.

2. How is cross-entropy loss computed? Roughly what is the loss when a 10-class network starts training?

Take the probability the network gives the correct answer and compute -log(probability). The closer the correct answer's probability is to 1, the closer the loss is to 0. At the start of training the 10 categories' probabilities are all around 0.1, so the loss is about -log(0.1) ≈ 2.30.

3. Why does the "loss of the last batch" swing up and down? How should you judge how training is going?

Each batch is only a small part of the data; drawing hard samples by chance makes the loss high, and the last batch has only a few images. Judge by metrics over the whole training and test sets, such as accuracy or the average loss over all the data, not the loss of one batch.