Build a Handwritten Digit Classifier

Train an MLP and CNN on MNIST to see why spatial structure matters.

Introduction

30 Second Summary

You recognize a handwritten 7 instantly, whether it's scribbled on a sticky note or neatly printed in a textbook. Teaching a computer to do the same thing forces you to think about what spatial structure in an image actually means.

In this project, you will build a handwritten digit classifier that evolves from a naive fully-connected network to a purpose-built CNN. You will train both architectures, watch the CNN surpass the baseline, and diagnose exactly where it still fails.

What You'll Build

You will feed thousands of handwritten digits through your trained CNN, see it correctly classify nearly every one, then open a gallery of its mistakes to understand exactly which spatial patterns still trip it up.

By the end of this project, you'll have:

  • A trained MLP baseline that you watch plateau at approximately 97% accuracy, proving firsthand that treating pixels as an unordered list loses critical spatial information.
  • A CNN you architect from scratch that jumps to approximately 99% accuracy by detecting local patterns like edges and curves regardless of their position in the image.
  • Training curves and a confusion matrix that let you compare both models side-by-side and pinpoint exactly which handwritten digits still fool your classifier.
  • Secret Mission: Apply regularization techniques like dropout and data augmentation to push past 99% accuracy and close the gap between training and test performance.

Do I need prior experience with PyTorch?

This project assumes you are comfortable with Python and have basic familiarity with PyTorch concepts like tensors and modules. You do not need prior experience building neural networks from scratch.

Before We Start

Before we dive in, let's lock in what you're about to build and why it matters. This is the moment you commit to the problem you're solving and think about how the structure of your input data shapes what a model can learn.

Set Up Your Deep Learning Environment

Deep learning projects pull in large libraries with specific version requirements. Installing them into your system Python can create conflicts with other projects on your machine.

A virtual environment gives this project its own isolated space for dependencies. In this step, you will create a dedicated folder, spin up a virtual environment, and install PyTorch (CPU-only) and matplotlib.

In this step, get ready to:
  • Create a project folder and Python virtual environment.
  • Install PyTorch (CPU-only) and matplotlib.
  • Verify that all libraries import successfully.
Create a project folder and Python virtual environment

A virtual environment is a self-contained copy of Python that keeps this project's packages separate from everything else on your system.

  • Press Cmd+Space (macOS) or the Windows key (Windows) to open your search bar.
  • Type Visual Studio Code and press Enter to open it.
  • Open the integrated terminal by pressing Ctrl+` (Windows) or Cmd+` (macOS).
  • Move to your Desktop by running this command:
cd ~/Desktop
  • Create your project folder and move into it by running these commands:
mkdir mnist-cnn
cd mnist-cnn
  • Create a Python virtual environment inside this folder by running:
python -m venv .venv

This creates a .venv folder containing a fresh, isolated Python installation.

  • Activate the virtual environment so that any packages you install go into it:

macOS

source .venv/bin/activate

Windows

.venv\Scripts\activate

You should see (.venv) appear at the start of your terminal prompt. This confirms the virtual environment is active.

Don't see (.venv) in your prompt?

  • On Windows, if you see an error about execution policies, run Set-ExecutionPolicy -Scope CurrentUser RemoteSigned in PowerShell, then try the activate command again.
  • Make sure you are inside the mnist-cnn folder. Run pwd to check your current directory.
  • help me activate my Python virtual environment
Install PyTorch (CPU-only) and matplotlib

PyTorch is the deep learning framework you will use to build and train your neural networks. The CPU-only version keeps the install small and avoids GPU driver complications.

  • Install PyTorch and its companion libraries by running this command:
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu

What are these three packages?

  • torch is the core PyTorch library for building neural networks and running tensor computations.
  • torchvision provides image datasets (including MNIST), transforms, and pre-built model architectures.
  • torchaudio is included as a standard companion package. You will not use it directly in this project.

This download may take a few minutes. You should see progress bars followed by a Successfully installed message at the end.

  • Install matplotlib for creating visualizations by running:
pip install matplotlib

You should see another Successfully installed message confirming matplotlib and its dependencies are in place.

Install command failed?

  • Check that (.venv) is still showing in your prompt. If it disappeared, re-activate the virtual environment before running pip.
  • If you see a connection error, check your internet connection and try running the command again.
  • If you see a version compatibility error, your Python version may be below 3.10. Run python --version to check. PyTorch requires Python 3.10 or later.
  • help me fix my PyTorch installation
Verify the installation

Before you run the verification command, what do you expect to see if everything installed correctly?

  • Confirm all libraries import successfully by running:
python -c "import torch; print(torch.__version__); import torchvision; import matplotlib; print('All imports OK')"

You should see the PyTorch version number printed on the first line, followed by All imports OK on the second line.

What does this command do?

The -c flag tells Python to execute the quoted string as a script. It imports each library and prints a confirmation. If any import fails, Python raises an error instead of printing the success message.

Seeing a ModuleNotFoundError?

  • Make sure your virtual environment is still active ((.venv) in the prompt). Packages were installed into the venv, so they are only visible when it is active.
  • If only matplotlib fails, re-run pip install matplotlib and check for errors in the output.
  • help me fix my import errors

Your deep learning environment is ready. Next up, you will load the MNIST handwritten digit dataset and explore what your model will be learning from.

Load MNIST and Explore the Data

Your environment is ready. Now it's time to meet your data.

Before building any model, you need to understand what you're working with. Understanding tensor shapes, pixel ranges, and the variation in handwriting styles builds the intuition you need for architectural decisions later. In this step, you'll download the MNIST dataset, visualize sample digits, and set up DataLoaders that feed batched data into your models.

In this step, get ready to:
  • Download MNIST and inspect tensor shapes.
  • Visualize a grid of sample handwritten digits.
  • Create DataLoaders for batched training.
Download MNIST and inspect the raw data

MNIST is a collection of 70,000 handwritten digit images. Each image is 28x28 pixels in grayscale. Your first job is to load this data and see what a single image looks like as a tensor.

  • Create a new file called explore_data.py in your mnist-cnn folder by running:
touch explore_data.py
  • Confirm the file exists by running:
ls

You should see explore_data.py listed alongside your .venv folder.

  • Open explore_data.py in VS Code and paste the following code to load MNIST and print its basic properties:
import torch
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt

# Define transforms: convert to tensor and normalize with MNIST statistics
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

# Load MNIST dataset
train_dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST('./data', train=False, download=True, transform=transform)

print(f"Training set size: {len(train_dataset)}")
print(f"Test set size: {len(test_dataset)}")
print(f"Single image shape: {train_dataset[0][0].shape}")
print(f"Label: {train_dataset[0][1]}")

What does this code do?

  • transforms.ToTensor() converts each image from a PIL image to a PyTorch tensor with pixel values between 0 and 1.
  • transforms.Normalize((0.1307,), (0.3081,)) standardizes pixel values using MNIST's precomputed mean and standard deviation. This helps the model train faster.
  • datasets.MNIST('./data', download=True, ...) downloads the dataset into a data/ folder (~50MB) if it doesn't already exist.
  • The four print statements reveal the dataset's key properties: number of samples, tensor dimensions, and the first label.
  • Save explore_data.py.
  • Run the script by executing:
python explore_data.py

You should see output like this:

Training set size: 60000
Test set size: 10000
Single image shape: torch.Size([1, 28, 28])
Label: 5

What do these shapes mean?

torch.Size([1, 28, 28]) means each image is 1 channel (grayscale), 28 pixels tall, and 28 pixels wide. The label is the digit's ground-truth value (0 through 9).

Seeing a download error?

  • Make sure you have an active internet connection. The first run downloads ~50MB of data.
  • Check that your virtual environment is still activated. You should see (.venv) at the start of your terminal prompt.
  • If you see a ModuleNotFoundError, re-run your install commands from Step 1.

Why can't Python find my torchvision module?

Visualize sample digits

Numbers alone don't build intuition. You need to see the handwriting variation: some 4s look like 9s, some 1s are slanted. This visual inspection helps you predict which digits might confuse a model later.

  • In explore_data.py, add the following code below the print statements to display a 4x4 grid of sample images:
# Visualize 16 sample images
fig, axes = plt.subplots(4, 4, figsize=(8, 8))
for i, ax in enumerate(axes.flat):
    image, label = train_dataset[i]
    ax.imshow(image.squeeze(), cmap='gray')
    ax.set_title(f"Label: {label}")
    ax.axis('off')
plt.tight_layout()
plt.savefig("sample_digits.png")
plt.show()

What does this code do?

  • plt.subplots(4, 4) creates a 4x4 grid of subplots to display 16 images at once.
  • image.squeeze() removes the channel dimension (from shape [1, 28, 28] to [28, 28]) so matplotlib can display it as a 2D grayscale image.
  • plt.savefig("sample_digits.png") saves the grid as a file so you can reference it later.
  • Save explore_data.py and run it again:
python explore_data.py

A matplotlib window pops up showing a 4x4 grid of handwritten digits. Each image has its ground-truth label printed above it. Close the window to continue.

What should I notice?

Look at the handwriting variation. Some digits are neatly centered. Others are slanted, thick, or thin. Pay attention to digits that look similar to each other. These are the pairs your model will struggle with later.

No window appearing?

  • If you're running in a remote or headless environment, plt.show() may not open a window. Check your project folder for sample_digits.png instead.
  • On Windows, make sure your terminal is not blocking the GUI window from appearing. Try clicking on the taskbar.

Why isn't matplotlib showing my plot window?

Create DataLoaders for batched training

Neural networks train on batches of images at a time, not one by one. A DataLoader wraps your dataset and serves up batches of a specified size, shuffling the training data each epoch to prevent the model from memorizing the order.

  • In explore_data.py, add the following code below the visualization block to create DataLoaders and confirm the batch shape:
# Create DataLoaders
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=1000, shuffle=False)

# Verify batch shape
images, labels = next(iter(train_loader))
print(f"Batch shape: {images.shape}")

What does this code do?

  • DataLoader(train_dataset, batch_size=64, shuffle=True) serves 64 images per batch, reshuffling the order each epoch so the model sees data in a different sequence every time.
  • The test loader uses batch_size=1000 and shuffle=False because evaluation doesn't need randomness and larger batches speed up inference.
  • next(iter(train_loader)) grabs the first batch so you can confirm the shape is what you expect.
  • Save explore_data.py.

Before you run the script, predict: what four numbers do you expect in the batch shape? Think about how many images per batch, the number of color channels, and the height and width of each image.

  • Run the script one more time:
python explore_data.py

After the dataset stats print and the grid window closes, you should see:

Batch shape: torch.Size([64, 1, 28, 28])

That's 64 images, each with 1 grayscale channel, 28 pixels tall, 28 pixels wide. This is the exact format your neural networks will consume in the next steps.

Shape doesn't match?

  • If you see a different batch size, double-check that your train_loader uses batch_size=64.
  • If you get an error on next(iter(train_loader)), make sure the DataLoader lines come after the dataset loading code.

Help me debug my DataLoader batch shape issue.

Where can I learn more about DataLoaders?

The PyTorch DataLoader documentation covers all the parameters including num_workers for parallel loading and pin_memory for GPU transfers.

✔️ Awesome, I've got everything!

Great. Double check you've saved explore_data.py.

ⓧ I'd like to double check the full code

Scroll up and confirm your explore_data.py contains these three sections in order:

  • Imports, transform definition, dataset loading, and four print statements.
  • The 4x4 visualization grid with plt.savefig("sample_digits.png") and plt.show().
  • DataLoader creation for train and test sets, followed by the batch shape print.

You now know exactly what your data looks like. You've seen the digit variation, confirmed tensor shapes, and verified that batches flow through the DataLoader correctly. Next up, you'll build your first neural network and see how far a naive approach can get on this data.

Build a Fully-Connected Baseline

In the last step, you explored MNIST digits and saw the variation in handwriting styles. Now it is time to build your first model and see how it performs.

The simplest approach is a fully-connected network (also called an MLP) that flattens each 28x28 image into a single list of 784 numbers. This destroys all spatial relationships between pixels. A pixel's position relative to its neighbors is completely lost. Training this baseline gives you a concrete performance ceiling to beat when you design the CNN later.

In this step, get ready to:
  • Define a fully-connected model that flattens 28x28 images into vectors.
  • Train and evaluate on MNIST to observe the performance ceiling.
  • Plot the training curve to see how quickly learning saturates.
Define the MLP architecture and training loop

You will create a new file called mlp_baseline.py that contains the model definition, training loop, and evaluation logic.

  • Create a new file called mlp_baseline.py in your mnist-cnn folder, then add the imports and data setup by pasting this code:
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt

# Data setup
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

train_dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST('./data', train=False, download=True, transform=transform)

train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=1000, shuffle=False)

What does this code do?

This is the same data pipeline you built in explore_data.py. It loads MNIST with the standard normalization (mean=0.1307, std=0.3081) and wraps the datasets in DataLoaders for batched training.

The extra imports (torch.nn, torch.optim) give you the building blocks for defining neural network layers and optimizers.

  • Below the data setup, add the MLP model class and instantiate it with the loss function and optimizer:

# MLP Model
class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.fc1 = nn.Linear(784, 128)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.flatten(x)
        x = self.relu(self.fc1(x))
        x = self.fc2(x)
        return x


model = MLP()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

What does this architecture do?

  • nn.Flatten() converts each 28x28 image into a flat vector of 784 numbers. This is the key decision that throws away spatial structure.
  • nn.Linear(784, 128) is a fully-connected layer mapping 784 inputs to 128 hidden neurons.
  • nn.ReLU() adds nonlinearity so the network can learn patterns beyond straight lines.
  • nn.Linear(128, 10) maps to 10 output classes (digits 0-9).
  • CrossEntropyLoss combines LogSoftmax and NLLLoss internally, so you pass raw logits directly.
  • Adam with lr=0.001 is a reliable default optimizer that adapts learning rates per parameter.
  • Save mlp_baseline.py and run it to confirm there are no syntax errors:
python mlp_baseline.py

The script should complete with no output and no errors. This confirms your imports and model definition are correct.

Seeing an ImportError?

Make sure your virtual environment is still activated. You should see (.venv) at the start of your terminal prompt. If not, reactivate it with .venv\Scripts\activate (Windows) or source .venv/bin/activate (macOS/Linux).

help me fix this ImportError in my MLP script

Now add the training loop. This is where the model actually learns from the data.

  • Below the optimizer line, add the training forward pass:
# Training loop
epochs = 5
train_losses = []

for epoch in range(epochs):
    model.train()
    running_loss = 0.0
    for images, labels in train_loader:
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()

    avg_loss = running_loss / len(train_loader)
    train_losses.append(avg_loss)

What is happening in this training loop?

Each epoch passes through all 60,000 training images in batches of 64. For each batch, the loop follows four steps of gradient descent:

  • optimizer.zero_grad() clears old gradients so they do not accumulate.
  • Forward pass: the images flow through the model to produce predictions.
  • loss.backward() computes how much each weight contributed to the error.
  • optimizer.step() nudges every weight in the direction that reduces the error.
  • Still inside the for epoch loop (indented one level), add the evaluation block below train_losses.append(avg_loss):
    # Evaluate
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in test_loader:
            outputs = model(images)
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    accuracy = 100 * correct / total
    print(f"Epoch {epoch+1}/{epochs} - Loss: {avg_loss:.4f} - Test Accuracy: {accuracy:.2f}%")

Why switch between train and eval mode?

model.train() enables training behaviors like dropout (if you had any). model.eval() disables them for consistent evaluation.

torch.no_grad() tells PyTorch not to track gradients during evaluation. This saves memory and speeds up inference since you are only measuring accuracy, not learning.

  • Save mlp_baseline.py.
Train and evaluate

Your training script is ready. Before you run it, consider this: you are feeding a model flat lists of 784 numbers with no spatial context. Each pixel is treated as an independent feature with no relationship to its neighbors.

  • Before you run the script, predict: what accuracy do you think a model that treats pixels as unordered numbers can reach on handwritten digits?
  • Run the training by executing:
python mlp_baseline.py

Each epoch takes 1-2 minutes on a typical laptop CPU. You will see output like this:

Epoch 1/5 - Loss: 0.3494 - Test Accuracy: 95.67%
Epoch 2/5 - Loss: 0.1567 - Test Accuracy: 96.58%
Epoch 3/5 - Loss: 0.1089 - Test Accuracy: 97.08%
Epoch 4/5 - Loss: 0.0830 - Test Accuracy: 97.15%
Epoch 5/5 - Loss: 0.0668 - Test Accuracy: 97.24%

Your exact numbers will differ slightly, but final test accuracy should land around 97%.

Why does this matter?

97% sounds impressive. But look closely: the loss drops quickly in the first two epochs, then barely improves. The network has learned everything it can from unstructured pixel values.

This is the performance ceiling you will break in the next step. The MLP cannot go higher because it has no way to detect local patterns like edges and curves. It sees each pixel in isolation.

Training seems stuck or accuracy is very low?

If accuracy stays below 90% after epoch 1, double-check that your forward method applies self.flatten(x) first. Without flattening, the Linear layer receives the wrong input shape and fails silently.

help me debug my MLP training that is not reaching 97% accuracy

Plot the training curve

A training curve visualizes how quickly the model learned and when it stopped improving. This plot makes the ceiling visceral.

  • At the bottom of mlp_baseline.py, add the plotting code:
# Plot training curve
plt.plot(range(1, epochs + 1), train_losses, marker='o', label='MLP Loss')
plt.xlabel('Epoch')
plt.ylabel('Training Loss')
plt.title('MLP Training Curve')
plt.legend()
plt.savefig("mlp_training_curve.png")
plt.show()

What does this plot show?

The curve plots training loss against epoch number. You should see a steep drop in the first 1-2 epochs followed by a flattening. This flat region is the model telling you: there is nothing more I can learn from pixel values alone.

  • Save mlp_baseline.py and run it again:
python mlp_baseline.py
  • After training completes, a matplotlib window opens showing your training curve. You should see the loss drop steeply at first, then flatten out.
  • Close the plot window. Confirm that mlp_training_curve.png was saved by running:
ls

You should see mlp_training_curve.png listed alongside your other files.

Plot window not appearing?

If you are running over SSH or in a headless environment, matplotlib may not be able to display the window. The plot is still saved as mlp_training_curve.png. Open it from your file explorer to verify.

help me get matplotlib plots to display on my system

✔️ Awesome, I've got everything!

Great. Double check you have saved mlp_baseline.py and that mlp_training_curve.png exists in your project folder.

ⓧ I'd like to double check the full code

Your mlp_baseline.py should contain these sections in order. Scroll up and confirm each one is present:

  • Imports (torch, torch.nn, torch.optim, torchvision, DataLoader, matplotlib).
  • Data setup (transform, train/test datasets, train/test DataLoaders).
  • MLP class with __init__ and forward methods.
  • Model instantiation, CrossEntropyLoss, and Adam optimizer.
  • Training loop (5 epochs) with inner batch loop and evaluation per epoch.
  • Plotting code that saves mlp_training_curve.png.

You now have a concrete number to beat. Your MLP flattens images, loses spatial information, and plateaus around 97%. In the next step, you will design a CNN architecture that preserves the spatial relationships between pixels and watch accuracy jump past this ceiling.

Design Your CNN Architecture

Your MLP reached about 97% accuracy and then stalled. The problem is structural: flattening a 28x28 image into a 784-element vector destroys all spatial relationships between pixels. A pixel knows nothing about its neighbors.

Convolutional layers fix this by sliding small filters across the image, detecting local patterns like edges and curves regardless of where they appear. In this step, you will redesign the network from scratch to preserve spatial structure and watch accuracy jump.

In this step, get ready to:
  • Design a CNN architecture with convolutional and pooling layers.
  • Train the CNN and compare its accuracy to the MLP baseline.
  • Save the trained model weights and plot the training curve.
Design and implement the CNN

Your CNN will live in a new file. The data loading and training infrastructure stays the same as the MLP. Only the model architecture changes.

  • Create a new file called cnn_model.py in your mnist-cnn folder and add the imports and data setup by pasting this code:
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt

# Data setup
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

train_dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST('./data', train=False, download=True, transform=transform)

train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=1000, shuffle=False)

What does this code do?

This is identical to the data setup in your MLP script. The same dataset, same transforms, same batch sizes. The only thing changing in this file is the model architecture.

Now for the key change. Instead of flattening the image into a vector, you will define a Convolutional Neural Network that preserves the 2D spatial structure.

  • Add the CNN class definition and model initialization below the data setup in cnn_model.py:


# CNN Model
class CNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.pool = nn.MaxPool2d(2)
        self.relu = nn.ReLU()
        self.flatten = nn.Flatten()
        self.fc1 = nn.Linear(64 * 7 * 7, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.pool(self.relu(self.conv1(x)))
        x = self.pool(self.relu(self.conv2(x)))
        x = self.flatten(x)
        x = self.relu(self.fc1(x))
        x = self.fc2(x)
        return x


model = CNN()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

What does this architecture do?

The CNN processes images through two convolutional blocks before passing to dense layers:

  • Conv2d(1, 32, kernel_size=3, padding=1) slides 32 small 3x3 filters across the image, detecting local features like edges and curves. padding=1 preserves the spatial dimensions so no information is lost at the borders.
  • MaxPool2d(2) halves the spatial dimensions (28 becomes 14, then 14 becomes 7). This provides translation invariance and reduces computation.
  • The second conv layer increases filters from 32 to 64, detecting more complex patterns composed of the simpler edges the first layer found.
  • After two pool operations, the image is 7x7 with 64 channels. Linear(64 * 7 * 7, 128) maps these spatial features to 128 dense units, then to 10 class outputs.

Why 3x3 kernels instead of larger ones?

A 3x3 kernel uses fewer parameters than a 5x5 kernel. Stacking two 3x3 layers gives the same receptive field as one 5x5, but with an extra nonlinearity between them. This is the pattern used in modern architectures like VGG and ResNet.

  • Save cnn_model.py and verify the architecture is valid by running:
python cnn_model.py
  • Your terminal prompt should return with no error messages and no output. This confirms PyTorch accepted your layer definitions and the model was instantiated successfully.

Seeing a TypeError or syntax error?

  • Check for typos in the layer names. kernel_size and padding must be spelled exactly.
  • Make sure class CNN(nn.Module): has the colon at the end and super().__init__() is in the __init__ method.
  • Confirm your virtual environment is activated. You should see (.venv) in your terminal prompt.

Help me fix my CNN class definition error.

Train the CNN and compare

The training loop is the same pattern you used for the MLP. The only difference is that the CNN receives 4D tensors (batch, channels, height, width) instead of flattened vectors.

  • Add the training loop below optimizer = optim.Adam(model.parameters(), lr=0.001) in cnn_model.py:

# Training loop
epochs = 5
train_losses = []

for epoch in range(epochs):
    model.train()
    running_loss = 0.0
    for images, labels in train_loader:
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()

    avg_loss = running_loss / len(train_loader)
    train_losses.append(avg_loss)

What does this code do?

This is identical to your MLP training loop. For each batch: zero the gradients, run the forward pass, compute the loss, backpropagate, and update weights. The CNN's convolutional layers participate in the same gradient descent process as the MLP's dense layers.

  • Still inside the for epoch loop, add the evaluation section below train_losses.append(avg_loss):

    # Evaluate
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in test_loader:
            outputs = model(images)
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    accuracy = 100 * correct / total
    print(f"Epoch {epoch+1}/{epochs} - Loss: {avg_loss:.4f} - Test Accuracy: {accuracy:.2f}%")

What does this code do?

After each training epoch, the model switches to evaluation mode with model.eval() and runs inference on the full test set without tracking gradients. This prints the test accuracy so you can watch it climb.

Before you run this, think: how much better do you expect the CNN to perform compared to the MLP's ~97%? Will the accuracy be slightly higher, significantly higher, or about the same?

  • Save cnn_model.py and run the training by executing:
python cnn_model.py

You should see accuracy climbing faster than the MLP did, reaching approximately 99% by epoch 5. Each epoch takes 1-2 minutes on CPU.

Compare: the MLP plateaued at ~97%. The CNN reaches ~99%. The architectural choice made the difference. Not more data, not longer training.

Training is extremely slow or accuracy stays low?

  • If training takes more than 3 minutes per epoch, check that you installed the CPU-only version of PyTorch. A misconfigured CUDA setup can cause silent slowdowns.
  • If accuracy stays below 90%, verify your forward method applies layers in the correct order: conv, relu, pool for each block.
  • Make sure the indentation of the evaluation section is inside the for epoch loop, not outside it.

Help me debug my CNN training loop.

Save the model and plot the training curve

Your trained CNN weights are valuable. Saving them means you can reload the model later without retraining.

  • Add the model saving and plotting code at the bottom of cnn_model.py (outside the for loop, at the top indentation level):

# Save model weights
torch.save(model.state_dict(), "mnist_cnn.pth")
print("Model saved to mnist_cnn.pth")

# Plot training curve
plt.plot(range(1, epochs + 1), train_losses, marker='o', label='CNN Loss')
plt.xlabel('Epoch')
plt.ylabel('Training Loss')
plt.title('CNN Training Curve')
plt.legend()
plt.savefig("cnn_training_curve.png")
plt.show()

What does this code do?

  • torch.save(model.state_dict(), "mnist_cnn.pth") saves only the learned weights (not the class definition). This is the recommended approach because it is portable and lightweight.
  • The matplotlib section plots loss vs. epoch, similar to your MLP plot. Comparing the two plots visually shows how much faster and lower the CNN's loss drops.
  • Save cnn_model.py and run it again:
python cnn_model.py
  • After training completes, you should see Model saved to mnist_cnn.pth printed in your terminal, followed by a matplotlib window showing the CNN loss curve.
  • Compare this curve to your mlp_training_curve.png from Step 3. The CNN loss should drop faster and reach a lower minimum.
  • Confirm the saved model file exists by running:
ls
  • You should see mnist_cnn.pth and cnn_training_curve.png alongside your other project files.

Don't see mnist_cnn.pth?

  • Make sure the torch.save line is outside the for loop (no indentation). If it is inside the loop, it would still work but run on every epoch.
  • Check that your script ran to completion without errors. If the training loop errored out, the save line was never reached.

Help me figure out why my model didn't save.

✔️ Awesome, I've got everything!

Great. Double check you have saved cnn_model.py and that mnist_cnn.pth exists in your project folder.

ⓧ I'd like to double check the full code

Scroll up and confirm your cnn_model.py contains each of these sections in order:

  • Imports (torch, nn, optim, datasets, transforms, DataLoader, matplotlib).
  • Data setup (transform, train/test datasets, train/test loaders).
  • CNN class with __init__ defining conv1, conv2, pool, relu, flatten, fc1, fc2, and a forward method chaining them.
  • Model instantiation (model = CNN()), criterion, and optimizer.
  • Training loop with 5 epochs, forward pass, backpropagation, and per-epoch evaluation printing accuracy.
  • Model save with torch.save and training curve plot with plt.savefig("cnn_training_curve.png").

You now have concrete proof that architecture matters. The same data, same loss function, same optimizer, same number of epochs. The only variable was whether the model preserves spatial structure. Next, you will interrogate this model to understand exactly where it still fails and why.

Diagnose and Understand Your Model

Your CNN hit ~99% accuracy in the last step. That sounds impressive, but it means roughly 100 test images are still being misclassified.

A model is not done when accuracy is high. It is done when you understand its failure modes. In this step, you will interrogate your trained CNN to find exactly which digits confuse it and why.

In this step, get ready to:
  • Generate predictions on the full test set and build a confusion matrix.
  • Visualize the misclassified images to understand failure patterns.
  • Identify the top confused digit pairs and reflect on what they reveal.
Generate predictions and build a confusion matrix

A confusion matrix is a 10x10 grid showing how often each true digit was predicted as each other digit. The diagonal shows correct predictions. Off-diagonal entries reveal exactly which pairs confuse the model.

To load your saved model weights, you need to redefine the CNN architecture first. The weights file only stores learned parameters, not the class structure.

  • Create a new file called diagnose.py in your mnist-cnn folder and add the imports and CNN class definition:
import torch
import torch.nn as nn
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
import numpy as np


# CNN Model (same definition as cnn_model.py)
class CNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.pool = nn.MaxPool2d(2)
        self.relu = nn.ReLU()
        self.flatten = nn.Flatten()
        self.fc1 = nn.Linear(64 * 7 * 7, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.pool(self.relu(self.conv1(x)))
        x = self.pool(self.relu(self.conv2(x)))
        x = self.flatten(x)
        x = self.relu(self.fc1(x))
        x = self.fc2(x)
        return x

Why redefine the CNN class?

When you saved the model with torch.save(model.state_dict(), ...), PyTorch stored only the learned weight values. To load them back, you need an identical architecture so PyTorch knows which weights go where.

The numpy import is new here. You will use it for array operations when building the confusion matrix.

  • Save diagnose.py and run it to confirm your imports and class definition are valid:
python diagnose.py

The script should finish without errors. No output is expected yet, but a crash-free run confirms all your imports resolved and the class definition is valid Python.

Seeing a ModuleNotFoundError?

Make sure your virtual environment is activated. You should see (.venv) in your terminal prompt.

If numpy is missing, it was installed automatically with PyTorch. Check that you are running from the mnist-cnn folder with the correct Python.

Why can't Python find my modules?

Now you will load your trained model from disk and run it against all 10,000 test images to collect predictions.

  • Add the data loading, model loading, and prediction collection code below your CNN class definition in diagnose.py:
# Load data
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])
test_dataset = datasets.MNIST('./data', train=False, download=True, transform=transform)
test_loader = DataLoader(test_dataset, batch_size=1000, shuffle=False)

# Load trained model
model = CNN()
model.load_state_dict(torch.load("mnist_cnn.pth", weights_only=True))
model.eval()

# Collect all predictions
all_preds = []
all_labels = []
all_images = []

with torch.no_grad():
    for images, labels in test_loader:
        outputs = model(images)
        _, predicted = torch.max(outputs, 1)
        all_preds.extend(predicted.numpy())
        all_labels.extend(labels.numpy())
        all_images.extend(images)

all_preds = np.array(all_preds)
all_labels = np.array(all_labels)

What does this code do?

  • model.load_state_dict(torch.load("mnist_cnn.pth", weights_only=True)) loads the weights you saved in Step 4 back into a fresh CNN instance.
  • weights_only=True is the safe default that avoids executing arbitrary code from pickle files.
  • model.eval() switches the model to evaluation mode, disabling training-specific behavior like dropout.
  • The torch.no_grad() context tells PyTorch not to track gradients, saving memory during inference.
  • Now add the confusion matrix construction and heatmap display below your prediction code:
# Build confusion matrix
confusion = np.zeros((10, 10), dtype=int)
for true, pred in zip(all_labels, all_preds):
    confusion[true][pred] += 1

# Display confusion matrix
fig, ax = plt.subplots(figsize=(10, 8))
im = ax.imshow(confusion, cmap='Blues')
ax.set_xlabel('Predicted')
ax.set_ylabel('True')
ax.set_title('Confusion Matrix')
ax.set_xticks(range(10))
ax.set_yticks(range(10))
for i in range(10):
    for j in range(10):
        ax.text(j, i, str(confusion[i][j]), ha='center', va='center',
                color='white' if confusion[i][j] > confusion.max() / 2 else 'black')
plt.colorbar(im)
plt.savefig("confusion_matrix.png")
plt.show()

What does this code do?

  • The 10x10 matrix counts how many times each true digit (row) was predicted as each digit (column). Correct predictions land on the diagonal.
  • ax.imshow(confusion, cmap='Blues') renders the matrix as a heatmap where darker cells mean higher counts.
  • The nested loop adds the numeric count to each cell so you can read exact values, not just color intensity.
  • Save diagnose.py and run it:
python diagnose.py

A heatmap window should appear showing your 10x10 confusion matrix. The diagonal will be deep blue (high correct counts). Look for any lighter-colored off-diagonal cells. Those reveal which digit pairs your CNN confuses.

Seeing a FileNotFoundError for mnist_cnn.pth?

Make sure you are running the script from inside the mnist-cnn folder where mnist_cnn.pth was saved in Step 4. Run ls and confirm the file is there.

Why can't torch.load find my model file?

Visualize misclassified examples

The confusion matrix shows you the numbers, but seeing the actual images tells you whether the errors are reasonable. Some misclassified digits are genuinely ambiguous, even to a human eye.

  • Add the misclassification analysis and gallery display at the bottom of diagnose.py:
# Find and display misclassified examples
misclassified_idx = np.where(all_preds != all_labels)[0]
print(f"\nTotal misclassified: {len(misclassified_idx)} out of {len(all_labels)}")

# Show top confused pairs
off_diagonal = confusion.copy()
np.fill_diagonal(off_diagonal, 0)
top_confused = np.unravel_index(np.argsort(off_diagonal.ravel())[-5:], (10, 10))
print("\nTop confused pairs (true -> predicted):")
for true, pred in zip(top_confused[0], top_confused[1]):
    print(f"  {true} -> {pred}: {confusion[true][pred]} times")

# Display misclassified images
fig, axes = plt.subplots(3, 5, figsize=(12, 8))
for i, ax in enumerate(axes.flat):
    if i < len(misclassified_idx):
        idx = misclassified_idx[i]
        image = all_images[idx].squeeze()
        ax.imshow(image, cmap='gray')
        ax.set_title(f"True: {all_labels[idx]}, Pred: {all_preds[idx]}")
        ax.axis('off')
plt.suptitle("Misclassified Examples")
plt.tight_layout()
plt.savefig("misclassified.png")
plt.show()

What does this code do?

  • np.where(all_preds != all_labels) finds every index where the model got it wrong.
  • The off-diagonal extraction removes correct predictions from the matrix, then sorts to find the 5 most-confused digit pairs.
  • The 3x5 gallery displays 15 misclassified images with their true and predicted labels so you can visually assess each failure.

Before you run this, think about which digit pairs you expect to be most confused. Which handwritten digits look most similar to each other?

  • Save diagnose.py and run it:
python diagnose.py

You should see the confusion matrix heatmap appear first (close it to continue), then the misclassified gallery. In your terminal, look for the printed statistics showing total misclassified count and the top confused pairs.

Gallery window not appearing?

Each plt.show() call blocks until you close the window. Close the confusion matrix window first, then the gallery will appear.

If no windows appear at all, your matplotlib backend may not support display. Check that you are running the script directly in a terminal, not inside a notebook.

Why aren't my matplotlib plots showing?

Summarize your findings

Look at the terminal output from your script. You should see something like:

Total misclassified: ~100 out of 10000

Top confused pairs (true -> predicted):
  4 -> 9: X times
  9 -> 4: X times
  3 -> 5: X times
  7 -> 2: X times
  5 -> 3: X times

What do these results tell you?

The errors concentrate on visually similar digits. A 4 with a closed top looks like a 9. A slanted 7 with a crossbar resembles a 2. A rounded 3 can look like a 5.

Look at the gallery images. Many of these are genuinely ambiguous, even to a human reader. Your model has learned the same intuitions about digit structure that you have. This is how practitioners debug models. The answer is usually in the data, not the architecture.

Your exact numbers will differ slightly because neural network training involves randomness. The pattern of confused pairs should be consistent though.

  • Confirm that confusion_matrix.png and misclassified.png were saved in your project folder by running:
ls

You should see both confusion_matrix.png and misclassified.png alongside your other project files.

✔️ Awesome, I've got everything!

Great. Make sure you have saved diagnose.py and both PNG files are in your project folder.

ⓧ I'd like to double check the full code

Scroll up and confirm your diagnose.py contains each of these sections in order:

  • Imports (torch, nn, datasets, transforms, DataLoader, matplotlib, numpy).
  • CNN class definition (identical to the one in cnn_model.py).
  • Data loading (test set only) and model loading from mnist_cnn.pth.
  • Prediction collection loop with torch.no_grad().
  • Confusion matrix build and heatmap display with plt.savefig("confusion_matrix.png").
  • Misclassified index finding, top confused pairs printing, and 3x5 gallery with plt.savefig("misclassified.png").

You now have a complete diagnostic toolkit for your CNN. You can see exactly where it succeeds and where it struggles. In the Secret Mission, you will use regularization and data augmentation to push past these limitations.

Secret mission

Add Regularization and Data Augmentation

Your CNN hits ~99% accuracy, but there is still a gap between training and test performance. Can you close it with dropout, batch normalization, and data augmentation to push past 99.2%?

Clean Up Your Resources

Clean Up Your Resources

This project runs entirely locally, so there are no ongoing costs. Decide whether to keep your resources running, pause them to come back later, or delete them entirely.

Resources you used:

  • The mnist-cnn project folder on your Desktop (all Python scripts, saved model weights, and generated plots).
  • The .venv virtual environment with PyTorch CPU and matplotlib installed.
  • The data/ folder containing the cached MNIST dataset (~50MB).

Keep everything running

No action needed. Choose this if you are still actively experimenting with your models or plan to extend the project soon.

  • Your virtual environment, trained model weights, and MNIST data remain on disk ready for immediate use.
  • You can rerun any script at any time by activating the venv and executing it with python.

Pause - I'll come back to this later

Shut down running processes to free up terminal resources. Your files and virtual environment stay on disk for later.

  • If your virtual environment is still active, deactivate it by running:
deactivate
  • When you want to come back, navigate to your project folder and reactivate the virtual environment.

macOS

cd ~/Desktop/mnist-cnn
source .venv/bin/activate

Windows

cd ~\Desktop\mnist-cnn
.venv\Scripts\activate

Delete - I don't want to use this again

Remove all project resources and start fresh. This deletes your scripts, trained model weights, MNIST data, and the virtual environment permanently.

  • If your virtual environment is still active, deactivate it first:
deactivate
  • Delete the entire mnist-cnn folder from your Desktop:

macOS

rm -rf ~/Desktop/mnist-cnn

Windows

rmdir /s /q %USERPROFILE%\Desktop\mnist-cnn
  • This removes all Python scripts, the .venv virtual environment, the data/ MNIST cache, saved plots, and the mnist_cnn.pth model weights.

Nice Work!

Nice Work!

Nice work! You just built a handwritten digit recognition system from the ground up, evolving it from a naive fully-connected network to a purpose-built CNN that classifies digits at 99% accuracy.

You've learned how to:

  • Built a fully-connected baseline that treats images as flat vectors, demonstrating firsthand why spatial structure matters for image classification.
  • Designed a CNN architecture from scratch with deliberate decisions about kernel size, filter count, pooling, and layer depth, then watched it surpass the MLP by ~2 percentage points.
  • Diagnosed model failures with a confusion matrix and misclassification gallery, building the debugging intuition needed for real deep learning work.
  • Secret Mission: Applied dropout, batch normalization, and data augmentation to push accuracy past 99.2% and close the generalization gap.

Ready to quiz yourself?