Project: Image Classifier End to End (Deep Learning Tutorial)

By the end of this project you will have built a complete image classifier twice. First, a NumPy digit reader that runs in your browser, with a proper train, validation and test split, Adam, early stopping, an error analysis and a test on a digit you draw yourself. Second, a PyTorch convolutional network for clothing photos that you run on your own computer, with sanity checks, a full training loop, a test report and an exported model.

This project uses almost every lesson in the course. Where a step needs more background, the link takes you back to it.

The plan

Every image classification project follows the same path, whatever the images are: digits, clothes, leaves from a farm near Nashik or products in a Jaipur warehouse.

Understand the data: how many images, which classes, how balanced, what they look like.
Split into training, validation and test sets.
Build a model that matches the data.
Check that the code can learn before training for long.
Train with a good optimiser, watching validation scores, keeping the best weights.
Test once on the untouched test set, and study the mistakes.
Save and export the model with everything needed to use it.

Part A: a digit reader in NumPy

The scikit-learn digits dataset has 1,797 handwritten digits, each an 8 × 8 grey image with pixel values from 0 to 16. Everything in Part A runs in your browser.

Step 1: understand the data

Before any model, look at the data. How many examples per class? What range do the pixel values have? What does one image look like?

PythonRun it yourself

import numpy as np
from sklearn.datasets import load_digits

X, y = load_digits(return_X_y=True)
print("images:", X.shape, " labels:", y.shape)
print("pixel range:", X.min(), "to", X.max())
print("examples per digit:", np.bincount(y))

for row in X[0].reshape(8, 8):                    # draw the first image with characters
    print("".join("#" if v > 8 else "+" if v > 3 else "." for v in row))
print("label:", y[0])

Output

images: (1797, 64)  labels: (1797,)
pixel range: 0.0 to 16.0
examples per digit: [178 182 177 183 181 182 181 179 174 180]
..+##...
..####+.
..#..#+.
.+#..++.
.++..#+.
.+#..#+.
..#+##..
..+##...
label: 0

The classes are balanced, with about 180 images of each digit, so accuracy is a fair headline metric. Pixels go from 0 to 16, so we will divide by 16 to scale them into 0 to 1, as you learned in A Two-Layer Network from Scratch.

Step 2: the full pipeline

The next program does steps 2 to 6 of the plan in one go. Read it section by section; each numbered comment matches a lesson you have already done.

Split: 20% test, then 20% of the rest for validation, all stratified so every digit appears in each part.
Model: 64 inputs, 64 ReLU hidden neurons, 10 outputs, with He initialisation (Weight Initialisation).
Training: mini-batches of 32 (Training Loops, Batches & Epochs), Adam (Optimisers) and a little weight decay.
Early stopping with a patience of 8 epochs, keeping the best weights (Overfitting & Regularisation).
Testing once at the end, with recall for each digit (Evaluating & Debugging Training).
Your own drawing: an 8 × 8 grid of # and . characters, turned into an image.

PythonRun it yourself

import numpy as np
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split

# 1. data: train / validation / test
X, y = load_digits(return_X_y=True)
X = X / 16.0
X_tmp, X_test, y_tmp, y_test = train_test_split(X, y, test_size=0.2, random_state=42, stratify=y)
X_train, X_val, y_train, y_val = train_test_split(X_tmp, y_tmp, test_size=0.2, random_state=42, stratify=y_tmp)

# 2. model: 64 -> 64 -> 10 with He initialisation
rng = np.random.default_rng(42)
params = {"W1": rng.normal(0, np.sqrt(2 / 64), (64, 64)), "b1": np.zeros(64),
          "W2": rng.normal(0, np.sqrt(2 / 64), (64, 10)), "b2": np.zeros(10)}

def forward(p, X):
    Z1 = X @ p["W1"] + p["b1"]
    A1 = np.maximum(0, Z1)
    Z2 = A1 @ p["W2"] + p["b2"]
    P = np.exp(Z2 - Z2.max(axis=1, keepdims=True))
    return Z1, A1, P / P.sum(axis=1, keepdims=True)

def loss_and_grads(p, X, y, weight_decay=1e-4):
    Z1, A1, P = forward(p, X)
    loss = -np.log(P[np.arange(len(y)), y] + 1e-12).mean()
    d = P.copy()
    d[np.arange(len(y)), y] -= 1
    d /= len(y)
    dZ1 = (d @ p["W2"].T) * (Z1 > 0)
    grads = {"W1": X.T @ dZ1 + weight_decay * p["W1"], "b1": dZ1.sum(axis=0),
             "W2": A1.T @ d + weight_decay * p["W2"], "b2": d.sum(axis=0)}
    return loss, grads

# 3. training with Adam, mini-batches and early stopping
m = {k: np.zeros_like(a) for k, a in params.items()}
v = {k: np.zeros_like(a) for k, a in params.items()}
t, best_acc, best, waited, best_epoch = 0, 0.0, None, 0, 0
for epoch in range(1, 61):
    order = rng.permutation(len(X_train))
    for start in range(0, len(order), 32):
        idx = order[start:start + 32]
        _, g = loss_and_grads(params, X_train[idx], y_train[idx])
        t += 1
        for k in params:
            m[k] = 0.9 * m[k] + 0.1 * g[k]
            v[k] = 0.999 * v[k] + 0.001 * g[k] ** 2
            params[k] -= 0.003 * (m[k] / (1 - 0.9 ** t)) / (np.sqrt(v[k] / (1 - 0.999 ** t)) + 1e-8)
    val_acc = np.mean(forward(params, X_val)[2].argmax(axis=1) == y_val)
    if val_acc > best_acc:
        best_acc, best, waited, best_epoch = val_acc, {k: a.copy() for k, a in params.items()}, 0, epoch
    else:
        waited += 1
    if epoch % 10 == 0 or waited == 8:
        print(f"epoch {epoch:>2}: validation accuracy {val_acc:.3f} (best {best_acc:.3f} at epoch {best_epoch})")
    if waited == 8:
        print("early stopping")
        break

# 4. final test with the best weights
params = best
pred = forward(params, X_test)[2].argmax(axis=1)
print(f"test accuracy: {np.mean(pred == y_test):.3f} on {len(y_test)} unseen digits")
recall = [np.mean(pred[y_test == c] == c) for c in range(10)]
worst = int(np.argmin(recall))
print("recall per digit:", np.round(recall, 2))
print(f"hardest digit: {worst}, often confused with {np.bincount(pred[(y_test == worst) & (pred != worst)], minlength=10).argmax()}")

# 5. try your own drawing: '#' is ink, '.' is paper (8 x 8)
drawing = """
..####..
.#....#.
......#.
.....#..
....#...
...#....
..#.....
.######.
"""
img = np.array([[1.0 if ch == "#" else 0.0 for ch in row] for row in drawing.split()]).reshape(1, 64)
probs = forward(params, img)[2][0]
top = probs.argsort()[::-1][:3]
print("your drawing looks like:", ", ".join(f"{d} ({probs[d]:.2f})" for d in top))

Output

epoch 10: validation accuracy 0.969 (best 0.972 at epoch 8)
epoch 16: validation accuracy 0.965 (best 0.972 at epoch 8)
early stopping
test accuracy: 0.950 on 360 unseen digits
recall per digit: [0.94 0.86 1.   0.97 1.   0.97 0.97 1.   0.8  0.97]
hardest digit: 8, often confused with 1
your drawing looks like: 2 (0.99), 3 (0.01), 8 (0.00)

Read the output from top to bottom:

Validation accuracy climbs quickly, then stops improving, and early stopping ends the run. The best epoch's weights are the ones kept.
The test accuracy is the honest score: these digits played no part in training or in choosing when to stop.
The recall line shows which digits are hardest, and the last line of the report says what the hardest one is mistaken for. That tells you where to look next.
The hand-drawn digit is read as a 2. Edit the drawing to make your own digits and see how confident the model is. Try a sloppy 8 or a 1 with a long base.

A note from building this project: we also tried data augmentation, shifting each training image by one pixel in a random direction. On these tiny 8 × 8 images it made the test accuracy clearly worse, because one pixel is a big move at this size. Augmentation is powerful, but it must suit the images. That is why we test every idea against the validation set instead of assuming it helps.

Part B: a clothing classifier in PyTorch

Now a bigger task: Fashion-MNIST, 70,000 grey 28 × 28 photos of clothing in 10 classes, such as T-shirt, trouser, sneaker and bag. This part is plain PyTorch code to run on your own computer, as set up in the introduction. Outputs shown are real where we ran the code on a laptop CPU; for the long training run, we tell you what to look for instead of giving numbers, because your results depend on your hardware and settings.

Step 3: data with a validation split and light augmentation

Clothing photos can be flipped left to right and still look like the same item, so a random horizontal flip is a sensible augmentation for training. Validation and test images get no randomness. (Datasets & DataLoaders explains each part.)

Python

import torch
from torch.utils.data import DataLoader, Subset
from torchvision import datasets
from torchvision.transforms import v2

torch.manual_seed(42)
base = [v2.ToImage(), v2.ToDtype(torch.float32, scale=True), v2.Normalize((0.5,), (0.5,))]
train_tf = v2.Compose([v2.RandomHorizontalFlip()] + base)
eval_tf = v2.Compose(base)

train_full = datasets.FashionMNIST("data", train=True, download=True, transform=train_tf)
val_full = datasets.FashionMNIST("data", train=True, download=True, transform=eval_tf)
test_set = datasets.FashionMNIST("data", train=False, download=True, transform=eval_tf)

order = torch.randperm(len(train_full), generator=torch.Generator().manual_seed(42)).tolist()
train_set, val_set = Subset(train_full, order[:50_000]), Subset(val_full, order[50_000:])
train_loader = DataLoader(train_set, batch_size=128, shuffle=True)
val_loader = DataLoader(val_set, batch_size=512)
test_loader = DataLoader(test_set, batch_size=512)
CLASSES = train_full.classes
print(len(train_set), len(val_set), len(test_set))
print(CLASSES)

Output

50000 10000 10000
['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat', 'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot']

Notice the trick for the split: two dataset objects over the same files, one with training transforms and one without, and the same shuffled index list cut into two parts. That way validation images never get flipped.

Step 4: the CNN

The model uses two convolution blocks (convolution, batch norm, ReLU, max pooling), as in Convolution & Pooling and Batch Normalisation, then dropout and a linear classifier.

Python

import torch
from torch import nn

class ClothesCNN(nn.Module):
    def __init__(self, n_classes=10):
        super().__init__()
        def block(c_in, c_out):
            return nn.Sequential(nn.Conv2d(c_in, c_out, 3, padding=1, bias=False),
                                 nn.BatchNorm2d(c_out), nn.ReLU(), nn.MaxPool2d(2))
        self.features = nn.Sequential(block(1, 32), block(32, 64))      # 28 -> 14 -> 7
        self.classifier = nn.Sequential(nn.Flatten(), nn.Dropout(0.3), nn.Linear(64 * 7 * 7, n_classes))

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

model = ClothesCNN()
x = torch.randn(4, 1, 28, 28)
print("features:", tuple(model.features(x).shape), "-> logits:", tuple(model(x).shape))
print(f"parameters: {sum(p.numel() for p in model.parameters()):,}")

Output

features: (4, 64, 7, 7) -> logits: (4, 10)
parameters: 50,282

The feature maps shrink from 28 × 28 to 7 × 7 while the channels grow to 64, and the whole model has only about 50,000 parameters: fewer than the 200,000 of the dense network in A Full PyTorch Training Loop, because convolutions share their weights.

Step 5: sanity checks before training

Before a long run, spend a few seconds on the two checks from Evaluating & Debugging Training: the starting loss should be near ln(10) = 2.30, and the model must be able to memorise one small batch. This uses ClothesCNN and train_set from the steps above.

Python

import math

images, labels = next(iter(DataLoader(train_set, batch_size=64, shuffle=False)))
torch.manual_seed(0)
model = ClothesCNN()
loss_fn = nn.CrossEntropyLoss()
print(f"starting loss {loss_fn(model(images), labels).item():.3f} vs ln(10) = {math.log(10):.3f}")

model.eval()                                    # no dropout while memorising one batch
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
for step in range(1, 101):
    loss = loss_fn(model(images), labels)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    if step in (1, 50, 100):
        print(f"step {step:>3}: loss on one batch {loss.item():.4f}")

Output

starting loss 2.587 vs ln(10) = 2.303
step   1: loss on one batch 2.3246
step  50: loss on one batch 0.0048
step 100: loss on one batch 0.0012

The starting loss is a little above 2.30, which is normal for a fresh model; a value of 10 or more would point to a problem. The loss on one fixed batch falls almost to zero within 50 steps. The data, model, loss and optimiser are wired correctly, so a long training run is worth the time. (We used model.eval() here only so that batch norm and dropout do not add noise to this quick check; real training uses model.train().)

Step 6: train with checkpoints and early stopping

This is the loop from A Full PyTorch Training Loop, with AdamW, a cosine learning rate schedule, a GPU if available (GPU Basics) and early stopping that saves a full checkpoint (Saving & Loading Models).

Python

device = torch.accelerator.current_accelerator().type if torch.accelerator.is_available() else "cpu"
torch.manual_seed(42)
model = ClothesCNN().to(device)
epochs = 15
optimizer = torch.optim.AdamW(model.parameters(), lr=2e-3, weight_decay=0.01)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
loss_fn = nn.CrossEntropyLoss()

def run_epoch(loader, train):
    model.train(train)
    total, correct, seen = 0.0, 0, 0
    with torch.set_grad_enabled(train):
        for images, labels in loader:
            images, labels = images.to(device), labels.to(device)
            logits = model(images)
            loss = loss_fn(logits, labels)
            if train:
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
            total += loss.item() * len(labels)
            correct += (logits.argmax(dim=1) == labels).sum().item()
            seen += len(labels)
    return total / seen, correct / seen

best_acc, patience, waited = 0.0, 4, 0
for epoch in range(1, epochs + 1):
    train_loss, train_acc = run_epoch(train_loader, train=True)
    val_loss, val_acc = run_epoch(val_loader, train=False)
    scheduler.step()
    print(f"epoch {epoch:>2}: train {train_loss:.3f}/{train_acc:.3f}  val {val_loss:.3f}/{val_acc:.3f}")
    if val_acc > best_acc:
        best_acc, waited = val_acc, 0
        torch.save({"model": model.state_dict(), "classes": CLASSES, "epoch": epoch,
                    "val_acc": val_acc, "normalise": (0.5, 0.5)}, "clothes_cnn.pt")
    else:
        waited += 1
        if waited >= patience:
            print("early stopping")
            break

What to look for while it runs:

Training and validation loss should both fall in the first epochs.
Validation accuracy should rise above the dense network's score from the training loop lesson, since a CNN suits images better. If it does not, recheck Step 5.
If training accuracy keeps climbing while validation accuracy stalls, the model has started to overfit and early stopping will soon end the run.

On a laptop CPU each epoch takes several minutes; on a GPU it is much faster. For a quick first try, set epochs = 3.

Step 7: test once and study the mistakes

Load the best checkpoint, score the test set once, and print a per-class report and the most common confusion.

Python

from sklearn.metrics import classification_report, confusion_matrix

ckpt = torch.load("clothes_cnn.pt", map_location=device)
model = ClothesCNN().to(device)
model.load_state_dict(ckpt["model"])
model.eval()

all_pred, all_true = [], []
with torch.inference_mode():
    for images, labels in test_loader:
        all_pred.append(model(images.to(device)).argmax(dim=1).cpu())
        all_true.append(labels)
pred, true = torch.cat(all_pred).numpy(), torch.cat(all_true).numpy()

print(f"best validation accuracy {ckpt['val_acc']:.3f} (epoch {ckpt['epoch']})")
print(classification_report(true, pred, target_names=ckpt["classes"], digits=3))
cm = confusion_matrix(true, pred)
cm[range(10), range(10)] = 0
i, j = divmod(cm.argmax(), 10)
print(f"most common mistake: {ckpt['classes'][i]} predicted as {ckpt['classes'][j]} ({cm[i, j]} times)")

In Fashion-MNIST the upper-body garments (T-shirt, shirt, pullover, coat) look alike at 28 × 28 pixels, so expect the "Shirt" class to have the lowest recall and most mistakes to fall among those four. Look at a few of the misclassified images: some are hard even for people.

Step 8: export for use in an app

Finally, export the trained model with torch.export so an app can load it without your training code, as described in Deployment. Save the class names and normalisation next to it; they are already in the checkpoint.

Python

model = model.cpu().eval()
example = torch.randn(1, 1, 28, 28)
exported = torch.export.export(model, (example,))
torch.export.save(exported, "clothes_cnn.pt2")

def predict(image_tensor):
    """image_tensor: (1, 28, 28) float tensor, already normalised like the training data."""
    with torch.inference_mode():
        probs = torch.softmax(model(image_tensor.unsqueeze(0)), dim=1)[0]
    best = int(probs.argmax())
    return CLASSES[best], round(float(probs[best]), 3)

Before sharing the model, write a short model card as in Responsible Deep Learning: what it is for, the data it was trained on, its test scores per class, and its limits (grey 28 × 28 product photos only, not real-world photos of people wearing clothes).

Ideas to extend the project

Add a third convolution block and compare validation accuracy and training time.
Replace the flatten and linear layer with global average pooling plus a smaller linear layer.
Try stronger augmentation such as v2.RandomAffine(degrees=10, translate=(0.1, 0.1)) and check whether it helps.
Fine-tune a pretrained ResNet-18 as in Transfer Learning, with grey images repeated to three channels.
Build your own small dataset (for example photos of five kinds of Indian snacks), organise it into one folder per class, and load it with torchvision.datasets.ImageFolder.

Common mistakes

Choosing settings by looking at the test set. Use the validation set for every decision; touch the test set once.
Random augmentation on validation or test images. Only training images get random transforms.
Skipping the sanity checks. Two minutes of checks can save hours of a broken run.
Losing the class names and preprocessing. Store them in the checkpoint, as Step 6 does.
Comparing runs with different seeds and splits. Fix the seed and the split so differences come from your change.

Key points

Every image project follows the same path: understand, split, build, check, train, test, export.
A NumPy network with Adam and early stopping reads 8 × 8 digits well, and you can test it on your own drawings.
A small CNN with batch norm and dropout suits 28 × 28 clothing photos better than a dense network, with fewer parameters.
Check the starting loss and overfit one batch before any long training run.
Save the best checkpoint with its metadata, test once, study the mistakes, export, and document the model.

Exercise

In Part A, change the hidden layer from 64 to 128 neurons and the Adam learning rate from 0.003 to 0.001. Run the program and compare the test accuracy, the stopping epoch and the hardest digit with the original. Then draw a digit 7 in the drawing grid and check the model's top three guesses.

Show solution

Change three places: the shapes in params, and the learning rate in the Adam update.

Python

params = {"W1": rng.normal(0, np.sqrt(2 / 64), (64, 128)), "b1": np.zeros(128),
          "W2": rng.normal(0, np.sqrt(2 / 128), (128, 10)), "b2": np.zeros(10)}
# ... and in the training loop:
#     params[k] -= 0.001 * (m[k] / (1 - 0.9 ** t)) / (np.sqrt(v[k] / (1 - 0.999 ** t)) + 1e-8)

drawing = """
.######.
......#.
.....#..
.....#..
....#...
....#...
...#....
...#....
"""

Note that the He scale for W2 uses the new fan-in, sqrt(2 / 128). With a smaller learning rate each epoch changes the weights less, so training usually needs more epochs before early stopping triggers; the wider layer gives more capacity. Compare your numbers with the original run instead of assuming the bigger model wins: on a dataset this small, the difference is often small either way.

Quick Quiz

1. In what order should the test set be used?

2. Why did random one-pixel shifts hurt the 8 × 8 digit model?

3. What should the starting loss of a 10-class classifier be, roughly?

4. Why does the CNN in Part B have fewer parameters than the dense network from the training loop lesson?

The plan

Part A: a digit reader in NumPy

Step 1: understand the data

Step 2: the full pipeline

Part B: a clothing classifier in PyTorch

Step 3: data with a validation split and light augmentation

Step 4: the CNN

Step 5: sanity checks before training

Step 6: train with checkpoints and early stopping

Step 7: test once and study the mistakes

Step 8: export for use in an app

Ideas to extend the project

Common mistakes

Key points

Exercise

Quick Quiz

Learn Deep Learning with a mentor