04 · Image Classification Project 🧪

5 min read

Time to put it all together. This is a complete, runnable image classifier that uses every skill from Phases 1–3 — your first real portfolio piece.

1. The 'Why'

Concepts learned in isolation don't stick until you've wired them together into something that actually runs. Every sub-module so far has handed you one puzzle piece — tensors, a model, a data pipeline, a training loop, checkpointing, convolutions — but the understanding that lasts comes from assembling those pieces into a working whole and watching a model learn end to end. This capstone does exactly that: we build a complete FashionMNIST classifier from data loading to final evaluation, and every step is something you've already met. Seeing the full pipeline in one place is what turns a collection of techniques into a mental template you can reach for on your next project.

We use FashionMNIST — 70,000 grayscale 28×28 images of clothing in 10 categories — because it's a built-in dataset (Module 04), trains in minutes on a laptop, yet is hard enough to be interesting (unlike plain digit MNIST). The plan is the standard workflow you'll reuse forever: load and split the data, build the CNN from Module 07, train it with the loop from Module 05, save the best model per Module 06, and finally report test accuracy with the metrics from Module 05. Follow it once here and you'll have both a runnable project to point to and a repeatable recipe for any classification task you meet in the wild.

2. Core Concepts

The end-to-end recipe. Data (Dataset/DataLoader + transforms) → model (CNN) → loss + optimizer → training loop with validation → best-model checkpoint → final test evaluation. Nothing new — just orchestration.

Device-agnostic code. Define one device and move the model and every batch to it (Modules 01, 05). The same script then runs on CPU, CUDA, or MPS.

Train/validation/test roles. Train fits the weights; validation picks the best epoch and guards against overfitting; test is the one-time honest score (Modules 04–05).

Reusing your helpers. The train_one_epoch and evaluate functions from Module 05 drop straight in — a sign your earlier code was structured well.

3. Code in Action

Data: load, split, batch

import torch
from torchvision import datasets, transforms
from torch.utils.data import DataLoader, random_split

device = "cuda" if torch.cuda.is_available() else "cpu"

tfms = transforms.Compose([
    transforms.ToTensor(),                          # -> (1, 28, 28) float in [0,1]
    transforms.Normalize((0.2860,), (0.3530,)),     # FashionMNIST mean/std
])

full_train = datasets.FashionMNIST("./data", train=True,  download=True, transform=tfms)
test_ds    = datasets.FashionMNIST("./data", train=False, download=True, transform=tfms)

# Carve a validation set out of the training data (reproducibly)
gen = torch.Generator().manual_seed(42)
train_ds, val_ds = random_split(full_train, [55000, 5000], generator=gen)

train_loader = DataLoader(train_ds, batch_size=64, shuffle=True)
val_loader   = DataLoader(val_ds,   batch_size=64, shuffle=False)
test_loader  = DataLoader(test_ds,  batch_size=64, shuffle=False)

Model, loss, optimizer

import torch.nn as nn

class SimpleCNN(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),  # -> 32x14x14
            nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2), # -> 64x7x7
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(64 * 7 * 7, 128), nn.ReLU(), nn.Dropout(0.25),
            nn.Linear(128, num_classes),
        )
    def forward(self, x):
        return self.classifier(self.features(x))

model = SimpleCNN().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

Train with validation + best-model saving, then test

import torch

def train_one_epoch(model, loader):
    model.train(); total = 0.0
    for X, y in loader:
        X, y = X.to(device), y.to(device)
        optimizer.zero_grad()
        loss = criterion(model(X), y)
        loss.backward(); optimizer.step()
        total += loss.item()
    return total / len(loader)

@torch.no_grad()
def evaluate(model, loader):
    model.eval(); correct = total = 0
    for X, y in loader:
        X, y = X.to(device), y.to(device)
        correct += (model(X).argmax(1) == y).sum().item()
        total += y.size(0)
    return correct / total

best_acc = 0.0
for epoch in range(10):
    tr_loss = train_one_epoch(model, train_loader)
    val_acc = evaluate(model, val_loader)
    if val_acc > best_acc:                         # keep the best (Module 06)
        best_acc = val_acc
        torch.save(model.state_dict(), "best_fashion_cnn.pt")
    print(f"epoch {epoch}: train_loss {tr_loss:.3f} | val_acc {val_acc:.3f}")

# Final, one-time test on the BEST model
model.load_state_dict(torch.load("best_fashion_cnn.pt", map_location=device))
print(f"Test accuracy: {evaluate(model, test_loader):.3f}")   # expect ~0.90+

4. Common Pitfalls

Pitfall 1 — Wrong normalization constants. FashionMNIST has its own mean/std (~0.286 / 0.353), not digit-MNIST's. Copy-pasting the wrong constants trains a slightly worse model. Use the dataset's own statistics (and the same ones at test time).

Pitfall 2 — Evaluating on the training set. It's tempting to report the high training accuracy, but that measures memorization. Report test accuracy from the held-out set, and only after model selection is done on validation.

Pitfall 3 — Leaving the model in train() mode for testing. With Dropout active, your test numbers become noisy and pessimistic. The evaluate helper calls model.eval() under @torch.no_grad() — keep that discipline (Modules 03, 05).

5. Further Reading & Watch List

⬅️ Prev: 03 · Transfer Learning · 🏁 Module complete! Back to the module index or on to Module 08 — NLP & Transformers.