Image Classifier
Train, evaluate, checkpoint, and diagnose an image classification pipeline in PyTorch
What You Will Build
Having built automatic differentiation (Project 4) and a NumPy neural network from scratch (Project 5), you now have the conceptual foundation to use PyTorch with complete understanding of what happens inside `loss.backward()` and `optimizer.step()`. In this project, you will build a reproducible computer vision training and evaluation pipeline on a classic benchmark such as MNIST or Fashion-MNIST. You will explore class distributions, normalize pixel intensities, wrap splits in `DataLoader` batches, compare a flattened MLP against a Convolutional Neural Network (`Conv2d`), track training and validation metrics, plot a 10×10 confusion matrix, inspect the highest-confidence misclassifications, and serialize/reload model checkpoints for standalone inference.
Why This Project Matters in the AI/LLM Stack
This project establishes the canonical PyTorch engineering workflow (`Dataset` -> `DataLoader` -> `nn.Module` -> `train/eval` loop -> `state_dict` checkpointing) that you will reuse for language modeling, Transformers, and LLM pretraining in Phases 4–8.
Learning Objectives
- Load, inspect, and normalize an image dataset and construct batched PyTorch `DataLoader` pipelines.
- Implement a clean `nn.Module` architecture (MLP and/or CNN) and separate `train_one_epoch` and `evaluate` functions.
- Generate a 10×10 confusion matrix and visually inspect misclassified images to diagnose failure patterns.
- Save model weights and metadata with `torch.save(state_dict)` and reload them into a fresh model instance for deterministic inference.
Prerequisites
- Project 5 (Neural Network From Scratch) and Phase 3 PyTorch fundamentals
- Understanding of `torch.Tensor`, `nn.Module`, `DataLoader`, and `model.train()` vs `model.eval()`
Required Mathematical Concepts
- Channel-wise pixel standardization: x_norm = (x - μ) / σ
- 2D spatial convolution and receptive fields: H_out = ⌊(H + 2P - K)/S⌋ + 1
- Multi-class Cross-Entropy loss over C = 10 class logits
- 10×10 Confusion Matrix where entry M_{i,j} counts true class i predicted as class j
Step-by-Step Implementation Tasks (12 Steps)
Check off each engineering task as your code passes its acceptance criteria. Progress is saved automatically in your browser.
Dataset Exploration, Normalization & DataLoaders
Load MNIST/Fashion-MNIST, inspect class balance and sample grids, normalize pixel tensors, and create DataLoaders.
MNIST (`28×28` grayscale handwritten digits `0–9`) or Fashion-MNIST provides 60,000 training images and 10,000 test images that train rapidly on CPU or GPU.
Checking label counts and plotting a grid of raw images ensures labels align with visual content.
Scaling uint8 pixels `[0, 255]` to zero-mean, unit-variance float32 tensors stabilizes gradient magnitudes in the first layer.
PyTorch `DataLoader` handles mini-batching and per-epoch shuffling automatically.
Network Architecture, Training Loop & Validation
Define the `nn.Module` classifier and implement reproducible training and `torch.no_grad()` validation loops.
Subclassing `torch.nn.Module` lets you define either a flattened MLP or a spatial CNN (`Conv2d -> ReLU -> MaxPool2d -> Linear`) that outputs 10 raw class logits.
Each training batch follows the five-step PyTorch rhythm: `model.train()`, forward pass, loss calculation, `optimizer.zero_grad()`, `loss.backward()`, and `optimizer.step()`.
During validation, we disable gradient tracking (`with torch.no_grad():`) and set `model.eval()` so Dropout and BatchNorm use evaluation behavior.
Accuracy Evaluation, Confusion Matrix & Error Analysis
Compute test accuracy, build a 10×10 confusion matrix, and visualize misclassified examples.
After selecting the best model checkpoint based on validation accuracy, final performance is measured once on the held-out 10,000-image test set.
A 10×10 confusion matrix reveals which digit pairs are most frequently confused (such as `4` vs `9` or `3` vs `5`).
Visualizing the model's top mistakes alongside their predicted softmax probabilities reveals ambiguous handwriting or systematic model blind spots.
Checkpoint Serialization & Standalone Inference
Save best checkpoints and verify reloading into a fresh model for single-image inference.
Saving `model.state_dict()` along with configuration metadata allows resuming training or deploying the model without retraining.
An inference function instantiates a fresh model, loads `state_dict`, sets `model.eval()`, and returns the predicted class and probability distribution for a single image.
Starter & Verification Code Examples
PyTorch Evaluation Loop & Checkpoint Save/Reload Pattern
Demonstrates the standard `@torch.no_grad()` evaluation function and `state_dict` checkpointing.
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
@torch.no_grad()
def evaluate(
model: nn.Module,
loader: DataLoader,
criterion: nn.Module,
device: torch.device,
) -> tuple[float, float, torch.Tensor]:
model.eval()
total_loss = 0.0
correct = 0
total = 0
confusion = torch.zeros(10, 10, dtype=torch.int64)
for images, labels in loader:
images, labels = images.to(device), labels.to(device)
logits = model(images)
loss = criterion(logits, labels)
total_loss += loss.item() * images.size(0)
preds = logits.argmax(dim=1)
correct += int((preds == labels).sum().item())
total += images.size(0)
for t, p in zip(labels.view(-1), preds.view(-1)):
confusion[int(t), int(p)] += 1
return total_loss / total, correct / total, confusionHow to Test Your Implementation
Test forward output shapes, single-batch overfitting sanity check, and checkpoint save/load determinism.
Single-Batch Overfit Sanity Check
Verify that the model and optimizer can drive training loss near zero on a single mini-batch of 16 images.
Steps: Take 1 batch of 16 images, run 50 optimizer steps, and check final accuracy on that batch.
Pass Criteria: Batch accuracy reaches 100% and loss drops below 0.02.
Checkpoint Save & Reload Logit Equivalence
Confirm that `state_dict` serialization preserves exact model behavior in `eval()` mode.
Steps: Save `model.state_dict()` to a temporary path, instantiate a fresh `DigitClassifier`, call `load_state_dict`, and compare outputs on a test batch.
Pass Criteria: `torch.allclose(out1, out2, atol=1e-6)` is True.
Overall Project Acceptance Criteria
- Dataset normalized using training statistics and wrapped in train/val/test DataLoaders.
- Training and `torch.no_grad()` validation loops implemented with `model.train()` and `model.eval()`.
- 10×10 confusion matrix and visual analysis of misclassified examples generated.
- Model checkpoint saved via `state_dict` and verified to produce identical inference outputs when reloaded.
Expandable Hints & Debugging Tips
Try implementing each stage on your own first. Expand a hint below if you want to verify tensor shapes, mathematical identities, or common edge cases.
Optional Stretch Challenges
MLP vs. CNN Parameter-Efficiency Benchmark
Compare a 2-layer MLP against a 2-layer CNN with fewer total parameters, measuring test accuracy and robustness to 2-pixel image shifts.
What You Learn: Why translation equivariance and weight sharing in convolutions outperform dense layers on spatial grids.
First-Layer Convolutional Filter & Feature Map Visualization
Plot the learned `3×3` or `5×5` convolutional kernels of `conv1` and the resulting activation maps on a test digit.
What You Learn: How early CNN layers learn edge, stroke, and corner detectors.
What Your Completed Repository Should Contain
A trained image classifier and reproducible evaluation report.
Deliverable Artifacts
- PyTorch image classification pipeline (`vision/model.py`, `vision/train.py`)
- 10×10 confusion matrix and top-misclassifications diagnostic plots
- Saved model checkpoint and standalone inference verification test
After Completing This Project, You Should Be Able To...
- Build end-to-end PyTorch data pipelines with normalization and batched `DataLoader` splits.
- Train and validate deep neural networks using proper `model.train()`, `model.eval()`, and `torch.no_grad()` semantics.
- Diagnose classification performance using per-class accuracy, confusion matrices, and visual error inspection.
- Serialize and reload trained model checkpoints (`state_dict`) for reproducible inference.
GitHub Repository Submission
Link your public or private GitHub repository URL to keep a structured portfolio of your from-scratch AI builds. Stored locally in your browser (self-reported; no GitHub tokens or remote code execution).
Suggested Repository: pytorch-image-classifier
Recommended Directory Layout:
vision/model.py vision/train.py vision/infer.py plots/confusion_matrix.png tests/test_image_classifier.py README.md
Recommended README Documentation Checklist
- Architecture summary and trainable parameter count
- Training/validation curves and 10×10 confusion matrix
- Misclassification analysis and inference usage instructions