Project 2 — MNIST: MLP vs CNN (Day 12)

Goal. Teach a network to read handwritten digits (0–9), train two designs on identical data, and compare them honestly. On MNIST a CNN usually beats a plain MLP even with a similar number of weights, but you are here to measure that, not to assume it.

What you will produce. Two trained models, a comparison table, loss curves, a confusion matrix, a gallery of mistakes, and a README. Why dropout, early stopping and convolutions work is taught in Day 12 (p. 235) and the training loop in Day 11 (p. 212); this walkthrough only runs them.

Before you start The code below uses PyTorch and has not been executed in the environment used to prepare this edition (it was syntax-checked, the parameter counts and tensor shapes were recomputed by hand). Run each step and compare with “You should see”. If something differs, the troubleshooting table at the end is the first place to look.

Step 1 — Imports, seed, device

COMPLETE RUNNABLE EXAMPLE Checked: syntax only, not run here

import copy
import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, random_split
from torchvision import datasets, transforms

torch.manual_seed(42)
np.random.seed(42)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)

The device line uses a GPU if you have one and falls back to the CPU. Everything in this project is small enough to run on a CPU, just slower. On a GPU, results can differ slightly between runs even with a seed.

Step 2 — Load MNIST and split it 55k / 5k / 10k

COMPLETE RUNNABLE EXAMPLE Checked: syntax only, not run here

tfm = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,)),     # MNIST's own pixel mean and std
])
full_train = datasets.MNIST("../data", train=True, download=True, transform=tfm)
test_set = datasets.MNIST("../data", train=False, download=True, transform=tfm)

train_set, val_set = random_split(
    full_train, [55000, 5000], generator=torch.Generator().manual_seed(42))

train_loader = DataLoader(train_set, batch_size=64, shuffle=True)
val_loader = DataLoader(val_set, batch_size=256)
test_loader = DataLoader(test_set, batch_size=256)
print(len(train_set), len(val_set), len(test_set))   # 55000 5000 10000

Why three sets: train to learn, validation to make choices (when to stop), test only at the very end. Normalize rescales pixels (0.1307 and 0.3081 are MNIST’s own mean and standard deviation) so training is faster and steadier. The first run downloads about 11 MB into data/ .

If it breaks: a download error means the machine cannot reach the MNIST host. Check your connection or proxy, or download the four MNIST files another way and place them in data/MNIST/raw/ .

Step 3 — Look at one batch

COMPLETE RUNNABLE EXAMPLE Checked: syntax only, not run here

X, y = next(iter(train_loader))
print(X.shape, y.shape)
for i in range(8):
    plt.subplot(2, 4, i + 1)
    plt.imshow(X[i].squeeze(), cmap="gray")
    plt.title(int(y[i]))
    plt.axis("off")
plt.show()

You should see: torch.Size([64, 1, 28, 28]) torch.Size([64]) , and eight digits. Read the shape as (batch, channels, height, width). Remember it: most PyTorch errors come from getting it wrong.

Step 4 — Build the MLP (flat pixels in, 10 scores out)

COMPLETE RUNNABLE EXAMPLE Checked: syntax only, not run here

class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Flatten(),                       # (B,1,28,28) -> (B,784)
            nn.Linear(784, 256), nn.ReLU(), nn.Dropout(0.2),
            nn.Linear(256, 128), nn.ReLU(), nn.Dropout(0.2),
            nn.Linear(128, 10))                 # 10 raw scores (logits)

    def forward(self, x):
        return self.net(x)

mlp = MLP()
print(sum(p.numel() for p in mlp.parameters()))

You should see: 235,146 parameters (784×256+256, 256×128+128, 128×10+10). Do not add a softmax at the end: CrossEntropyLoss already includes it and expects raw scores.

Step 5 — Write the training helpers once, reuse for both

models

These two functions are the heart of the project. Type them slowly, and keep the indentation exactly as shown: the three hist[...] lines belong inside the epoch loop.

COMPLETE RUNNABLE EXAMPLE Checked: syntax only, not run here

def evaluate(model, loader, loss_fn):
    model.eval()                                # turns dropout OFF
    total, correct, n = 0.0, 0, 0
    with torch.no_grad():                       # no learning here
        for X, y in loader:
            X, y = X.to(device), y.to(device)
            out = model(X)
            total += loss_fn(out, y).item() * len(y)
            correct += (out.argmax(1) == y).sum().item()
            n += len(y)
    return total / n, correct / n

def fit(model, max_epochs=30, patience=3, lr=1e-3):
    model.to(device)
    loss_fn = nn.CrossEntropyLoss()
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    hist = {"train_loss": [], "val_loss": [], "val_acc": []}
    best_val, best_state, bad = float("inf"), None, 0
    for epoch in range(1, max_epochs + 1):

(continued)

    model.train()                           # turns dropout ON
    run, n = 0.0, 0
    for X, y in train_loader:
        X, y = X.to(device), y.to(device)
        opt.zero_grad()                     # 1. clear old gradients
        loss = loss_fn(model(X), y)         # 2. forward + loss
        loss.backward()                     # 3. backpropagation
        opt.step()                          # 4. update weights
        run += loss.item() * len(y)
        n += len(y)
    vl, va = evaluate(model, val_loader, loss_fn)
    hist["train_loss"].append(run / n)
    hist["val_loss"].append(vl)
    hist["val_acc"].append(va)
    print(f"epoch {epoch:2d}  train {run / n:.4f}  val {vl:.4f}  val_acc {va:.4f}")
    if vl < best_val:
        best_val, bad = vl, 0
        best_state = copy.deepcopy(model.state_dict())   # remember the best weights
    else:
        bad += 1
        if bad >= patience:
            print("early stopping")
            break
if best_state is not None:
    model.load_state_dict(best_state)       # go back to the best epoch
return hist

The learning loop in plain words: clear old corrections, make predictions and measure the error, work out which weights caused it, nudge those weights. Early stopping stops when validation loss has not improved for 3 epochs in a row, then restores the best weights. The validation set picks the stopping point; the test set is never involved.

Step 6 — Train and test the MLP

COMPLETE RUNNABLE EXAMPLE Checked: syntax only, not run here

loss_fn = nn.CrossEntropyLoss()
mlp_hist = fit(mlp)
mlp_test_loss, mlp_test_acc = evaluate(mlp, test_loader, loss_fn)
print("MLP test accuracy:", mlp_test_acc)

You should see: validation accuracy climbing past 97% within a few epochs and a test accuracy around 97–98% (typical). Validation loss rising while training loss keeps falling is overfitting, and early stopping handles it.

Step 7 — Build the CNN

COMPLETE RUNNABLE EXAMPLE Checked: syntax only, not run here

class CNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 16, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),    # 28 -> 14
            nn.Conv2d(16, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2))   # 14 -> 7
        self.classifier = nn.Sequential(
            nn.Flatten(),                       # 32 x 7 x 7 = 1568
            nn.Linear(32 * 7 * 7, 128), nn.ReLU(), nn.Dropout(0.3),
            nn.Linear(128, 10))

    def forward(self, x):
        return self.classifier(self.features(x))

cnn = CNN()
print(cnn(torch.randn(2, 1, 28, 28)).shape)     # shape check: torch.Size([2, 10])
print(sum(p.numel() for p in cnn.parameters()))

You should see: torch.Size([2, 10]) and 206,922 parameters (the MLP has 235,146, so the comparison is roughly like for like). Always run a shape check with random numbers before training. It catches the classic “mat1 and mat2 shapes cannot be multiplied” error in one second. If it fails, print cnn.features(torch.randn(2,1,28,28)).shape and fix the number inside nn.Linear(...) .

Step 8 — Train the CNN the same way

COMPLETE RUNNABLE EXAMPLE Checked: syntax only, not run here

cnn_hist = fit(cnn)
cnn_test_loss, cnn_test_acc = evaluate(cnn, test_loader, loss_fn)
print("CNN test accuracy:", cnn_test_acc)

You should see: typically about 98.5–99.3% test accuracy. The same training function and the same data make this a fair comparison. If your CNN does not beat your MLP, that is a result to report and investigate (learning rate, early stopping, seed), not a bug by itself.

Step 9 — Fill the comparison table

Instead of filling the table by hand, let the code do it, so the numbers in your README are the numbers you actually got.

COMPLETE RUNNABLE EXAMPLE Checked: syntax only, not run here

import pandas as pd

n_params = lambda m: sum(p.numel() for p in m.parameters())
table = pd.DataFrame({
    "model": ["MLP", "CNN"],
    "parameters": [n_params(mlp), n_params(cnn)],
    "best_val_acc": [max(mlp_hist["val_acc"]), max(cnn_hist["val_acc"])],
    "test_acc": [mlp_test_acc, cnn_test_acc],
    "epochs_run": [len(mlp_hist["val_loss"]), len(cnn_hist["val_loss"])],
}).round(4)
print(table)

Step 10 — Make the three plots

The test set is used here only to look at the finished CNN. Nothing in these plots may change a model, a setting or a stopping point.

COMPLETE RUNNABLE EXAMPLE Checked: syntax only, not run here

from sklearn.metrics import ConfusionMatrixDisplay

# 1. Loss curves
for name, h in [("MLP", mlp_hist), ("CNN", cnn_hist)]:
    plt.plot(h["train_loss"], label=f"{name} train")
    plt.plot(h["val_loss"], "--", label=f"{name} val")
plt.xlabel("epoch")
plt.ylabel("loss")
plt.title("Train vs validation loss")
plt.legend()
plt.show()

# 2. Confusion matrix for the CNN on the test set
cnn.eval()
preds, trues, imgs = [], [], []
with torch.no_grad():
    for X, y in test_loader:
        preds.append(cnn(X.to(device)).argmax(1).cpu())
        trues.append(y)
        imgs.append(X)
preds, trues, imgs = torch.cat(preds), torch.cat(trues), torch.cat(imgs)
ConfusionMatrixDisplay.from_predictions(trues, preds)
plt.title("CNN confusion matrix (test set)")
plt.show()

# 3. Eight mistakes
wrong = (preds != trues).nonzero().flatten()[:8]
for i, idx in enumerate(wrong):
    plt.subplot(2, 4, i + 1)
    plt.imshow(imgs[idx].squeeze(), cmap="gray")

(continued)

    plt.title(f"true {trues[idx].item()} pred {preds[idx].item()}")
    plt.axis("off")
plt.show()

How to read them: in the loss plot, the gap between train and validation shows overfitting. In the confusion matrix, bright cells off the diagonal are common mix-ups (often 4 with 9, and 3 with 5). The mistake gallery usually shows digits that a human would also find sloppy.

Step 11 — Save and write up

COMPLETE RUNNABLE EXAMPLE Checked: syntax only, not run here

import os

os.makedirs("../outputs", exist_ok=True)
torch.save(cnn.state_dict(), "../outputs/cnn_mnist.pt")

Saving state_dict() (the weights only) is the recommended way. In your README, write three conclusions in your own words: (1) why the CNN did or did not win with similar parameters (it keeps the image’s 2-D structure and reuses the same small filters everywhere), (2) what dropout and early stopping did to your curves, (3) which digits get confused and why.

You are done when…

n Both models use the same data splits and the same fit function. n You printed parameter counts for both (235,146 and 206,922). n Validation data chose the stopping point; the test set was used once per model, after training ended. n You have the comparison table, three plots and a README with three conclusions.

If you get stuck

Error or symptom Fix

mat1 and mat2 shapes cannot be Wrong number in the first nn.Linear after flatten. Print the multiplied feature shape (Step 7).

Expected 4D input for Conv2d Input must be (B, C, H, W). Add a channel with

x.unsqueeze(1)

CrossEntropyLoss complains about Labels must be integer class numbers of shape (B,), not the target one-hot.

IndentationError in fit The hist[...] lines or plt.show() lost their indentation. Compare with the code above.

Validation accuracy jumps around You forgot model.eval() before validating (it is inside

evaluate

Training is very slow Use a GPU, or raise the batch size to 128.

Results get worse after early You forgot to restore the best weights. It is already inside stopping fit .

Download fails See Step 2: connection, proxy, or place the files in

data/MNIST/raw/