Digit Recognizer: A First Real CNN

What We're Building

A digit recognizer for MNIST — 70,000 real handwritten digit images (0-9), the classic "hello world" of computer vision, comparable in stature to Iris for classification or IMDb for sentiment. It ships directly through torchvision, needs no cleaning, and is small enough to train in a reasonable time even without a GPU.

Step 1: Loading the Real Dataset

import torch
from torchvision import datasets, transforms

transform = transforms.ToTensor()
train_data = datasets.MNIST(root="data", train=True, download=True, transform=transform)
test_data = datasets.MNIST(root="data", train=False, download=True, transform=transform)

print(len(train_data), len(test_data))
print(train_data[0][0].shape)   # torch.Size([1, 28, 28]) — 1 grayscale channel, 28x28 pixels

Step 2: Looking at the Actual Images

import matplotlib.pyplot as plt

fig, axes = plt.subplots(1, 5, figsize=(10, 2))
for i, ax in enumerate(axes):
    image, label = train_data[i]
    ax.imshow(image.squeeze(), cmap="gray")
    ax.set_title(str(label))
    ax.axis("off")
plt.savefig("sample_digits.png")

Worth seeing before building anything — these are genuinely messy, human-written digits, not clean printed text. A model that works here has to tolerate real handwriting variation, not just recognize one fixed font.

Step 3: The CNN — nn1-7's Own Architecture, Applied

import torch.nn as nn

class DigitCNN(nn.Module):
    def __init__(self):
        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.pool = nn.MaxPool2d(2)
        self.fc = nn.Linear(32 * 7 * 7, 10)

    def forward(self, x):
        x = self.pool(torch.relu(self.conv1(x)))   # 28x28 -> 14x14
        x = self.pool(torch.relu(self.conv2(x)))   # 14x14 -> 7x7
        x = x.view(x.size(0), -1)             # flatten for the final linear layer
        return self.fc(x)

Per nn1-7: each Conv2d layer slides a small learned kernel across the image, extracting local features (edges, curves, strokes) rather than treating each pixel independently. MaxPool2d(2) halves the spatial size after each convolution, keeping the strongest signal from each small region — exactly the convolution-then-pooling pattern that chapter covered, applied here to real pixels instead of illustrative diagrams.

Step 4: Training — Batching, Reused From Chapter 1

from torch.utils.data import DataLoader

train_loader = DataLoader(train_data, batch_size=64, shuffle=True)
test_loader = DataLoader(test_data, batch_size=64)

model = DigitCNN()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
loss_fn = nn.CrossEntropyLoss()

for epoch in range(3):
    for images, labels in train_loader:
        optimizer.zero_grad()
        loss = loss_fn(model(images), labels)
        loss.backward()
        optimizer.step()
    print(f"Epoch {epoch+1} done.")

The same DataLoader-based batching pattern dsproj2-1 already introduced — no new concept here, just applied to image tensors instead of encoded text sequences.

Step 5: Evaluating

correct, total = 0, 0
with torch.no_grad():
    for images, labels in test_loader:
        preds = model(images).argmax(dim=1)
        correct += (preds == labels).sum().item()
        total += labels.size(0)

print(f"Test accuracy: {correct / total:.2%}")   # typically well above 98%

Step 6: Looking at What It Actually Gets Wrong

wrong = []
with torch.no_grad():
    for images, labels in test_loader:
        preds = model(images).argmax(dim=1)
        mismatches = (preds != labels).nonzero()
        for idx in mismatches[:5]:
            i = idx.item()
            wrong.append((images[i], labels[i].item(), preds[i].item()))
        if wrong:
            break

fig, axes = plt.subplots(1, len(wrong), figsize=(10, 2))
for ax, (img, actual, predicted) in zip(axes, wrong):
    ax.imshow(img.squeeze(), cmap="gray")
    ax.set_title(f"actual {actual}, predicted {predicted}")
    ax.axis("off")
plt.savefig("misclassified.png")

What's Next

Chapter 5: Weather Forecaster — returning to dsproj1-3's own logged weather data, this time treating it as a genuinely different task shape: predicting tomorrow from today, not classifying independent rows.