Capstone: Fine-Tune and Ship an Image Classifier

Project brief

You are the ML engineer at a pet-insurance startup. Customers upload a photo of their cat or dog when they file a claim, and the claims team wants the breed pre-filled so they can price risk and catch obviously mismatched policies. Your job is to ship a breed classifier that runs on a CPU server, report honestly how good it is, and explain where it fails. This is the most common shape of real computer-vision work: nobody trains from scratch; you take a network pretrained on millions of images, adapt it to a few thousand labelled examples, measure it carefully, and package it so another team can run it.

By the end you will have a fine-tuned model, a confusion matrix and an error analysis, an exported ONNX file that gives the same answers as your PyTorch model, and a short technical report.

Learning goals

• Build a reproducible image pipeline: deterministic train/validation/test splits, train-time augmentation, and eval-time preprocessing that matches the pretrained weights.
• Fine-tune a pretrained CNN (ResNet-50) and optionally a Vision Transformer, comparing a linear probe against full fine-tuning.
• Write a correct training loop with mixed precision, learning-rate scheduling, checkpointing on validation accuracy, and fixed random seeds.
• Evaluate with accuracy, per-class precision and recall, macro-F1 and a confusion matrix, and do a structured error analysis.
• Export to ONNX and prove numerical equivalence with ONNX Runtime.

Required background

Complete or review these first: convolutional-neural-networks (what convolutions, pooling and residual blocks do), training-deep-networks (optimizers, learning-rate schedules, regularization), model-evaluation-and-metrics (precision, recall, F1, confusion matrices), computer-vision (transfer learning in vision) and, for the export milestone, deploying-models.

Setup

Environment. Python 3.10 or newer with torch, torchvision, scikit-learn, matplotlib, onnx and onnxruntime. Recent PyTorch releases use a newer ONNX exporter that may also need the onnxscript package; install it if export asks for it. Pin the versions you used in a requirements.txt and record them in your report.

Hardware. A single GPU with 8 GB or more of memory trains ResNet-50 at batch size 32 comfortably in mixed precision. If you have no GPU, free hosted notebook services (Google Colab and Kaggle Notebooks are the best known) give time-limited access to a single entry-level GPU; quotas and the GPU model change, so save checkpoints to persistent storage every epoch. Apple-silicon Macs can use PyTorch's mps device. A CPU alone can do the linear-probe milestone overnight but is too slow for full fine-tuning.

Dataset. The Oxford-IIIT Pet Dataset (Parkhi, Vedaldi, Zisserman and Jawahar, CVPR 2012) from the Visual Geometry Group at Oxford: 37 breeds (12 cats, 25 dogs), roughly 200 images per breed, 7,349 images in total, with an official trainval split of 3,680 images and a test split of 3,669. Home page: https://www.robots.ox.ac.uk/~vgg/data/pets/. The dataset page states a Creative Commons Attribution-ShareAlike 4.0 licence; check it yourself before any commercial use, cite the paper in your report, and note that some images were collected from the web, so copyright in individual photos may sit with their owners. torchvision.datasets.OxfordIIITPet downloads it for you.

Memory check. ResNet-50 has about 25.6 million parameters, so its FP32 weights are about $25.6\text{M} \times 4\ \text{bytes} \approx 102$ MB. Training also stores gradients (another 102 MB) and two AdamW moment buffers (204 MB), about 410 MB before activations. Activations dominate: they scale with batch size and image resolution, which is why batch size is the first knob to turn when you run out of memory.

Milestone 1: Data pipeline and baseline sanity checks

Build the splits once and save the indices to disk so every experiment uses the same data. Hold out 15% of trainval as validation; the official test split is touched only at the very end.

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

ROOT, SEED = "data", 0
MEAN, STD = (0.485, 0.456, 0.406), (0.229, 0.224, 0.225)  # ImageNet stats

train_tf = transforms.Compose([
    transforms.RandomResizedCrop(224, scale=(0.6, 1.0)),
    transforms.RandomHorizontalFlip(),
    transforms.ColorJitter(0.2, 0.2, 0.2),
    transforms.ToTensor(), transforms.Normalize(MEAN, STD)])
eval_tf = transforms.Compose([
    transforms.Resize(256), transforms.CenterCrop(224),
    transforms.ToTensor(), transforms.Normalize(MEAN, STD)])

base = datasets.OxfordIIITPet(ROOT, split="trainval", download=True)
g = torch.Generator().manual_seed(SEED)
perm = torch.randperm(len(base), generator=g).tolist()
n_val = int(0.15 * len(base))
split = {"val": perm[:n_val], "train": perm[n_val:]}
json.dump(split, open("split.json", "w"))

train_ds = Subset(datasets.OxfordIIITPet(ROOT, "trainval", transform=train_tf), split["train"])
val_ds = Subset(datasets.OxfordIIITPet(ROOT, "trainval", transform=eval_tf), split["val"])
test_ds = datasets.OxfordIIITPet(ROOT, "test", transform=eval_tf, download=True)
classes = base.classes
train_dl = DataLoader(train_ds, batch_size=32, shuffle=True, num_workers=2)
val_dl = DataLoader(val_ds, batch_size=64, num_workers=2)
test_dl = DataLoader(test_ds, batch_size=64, num_workers=2)
print(len(train_ds), len(val_ds), len(test_ds), len(classes))

Two separate dataset objects are created over the same files so that validation images get the deterministic eval_tf while training images get random augmentation; wrapping one dataset and changing its transform later is a classic bug that silently augments your validation set. Before training, plot a grid of 16 augmented training images with their labels, check the class counts are roughly balanced, and compute the majority-class baseline (about 1/37, or 2.7%). Any trained model must beat that by a wide margin.

Milestone 2: Linear probe, then full fine-tuning

A linear probe freezes the pretrained backbone and trains only a new final layer. It is fast, hard to overfit, and tells you how good the pretrained features already are. Full fine-tuning then unfreezes everything with a smaller learning rate for the backbone than for the new head (discriminative learning rates), because the pretrained weights are already close to useful and large updates would destroy them.

import torch.nn as nn
from torchvision import models

def build_model(num_classes=37, freeze_backbone=True):
    weights = models.ResNet50_Weights.DEFAULT
    model = models.resnet50(weights=weights)
    for p in model.parameters():
        p.requires_grad = not freeze_backbone
    model.fc = nn.Linear(model.fc.in_features, num_classes)  # new head, trainable
    return model

def make_optimizer(model, head_lr=1e-3, backbone_lr=1e-4, wd=0.05):
    head = list(model.fc.parameters())
    head_ids = {id(p) for p in head}
    backbone = [p for p in model.parameters()
                if p.requires_grad and id(p) not in head_ids]
    groups = [{"params": head, "lr": head_lr}]
    if backbone:
        groups.append({"params": backbone, "lr": backbone_lr})
    return torch.optim.AdamW(groups, weight_decay=wd)

Run the probe for about 5 epochs, then full fine-tuning for 10 to 15 epochs starting from the ImageNet weights. Optional comparison: swap in models.vit_b_16(weights=models.ViT_B_16_Weights.DEFAULT) and replace model.heads.head; ViT-B/16 has about 86 million parameters, so expect slower epochs and a stronger dependence on augmentation and learning-rate warmup.

Milestone 3: A correct training loop

import time

def run_epoch(model, loader, device, optimizer=None, scheduler=None, scaler=None):
    train = optimizer is not None
    model.train(train)
    loss_fn = nn.CrossEntropyLoss(label_smoothing=0.1 if train else 0.0)
    total, correct, loss_sum = 0, 0, 0.0
    for x, y in loader:
        x, y = x.to(device), y.to(device)
        with torch.set_grad_enabled(train), torch.autocast(
                device_type=device, dtype=torch.float16, enabled=(device == "cuda")):
            logits = model(x)
            loss = loss_fn(logits, y)
        if train:
            optimizer.zero_grad(set_to_none=True)
            scaler.scale(loss).backward()
            scaler.step(optimizer); scaler.update(); scheduler.step()
        loss_sum += loss.item() * len(y)
        correct += (logits.argmax(1) == y).sum().item(); total += len(y)
    return loss_sum / total, correct / total

def fit(model, epochs, device):
    torch.manual_seed(SEED)
    model.to(device)
    opt = make_optimizer(model)
    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=epochs * len(train_dl))
    scaler = torch.amp.GradScaler(enabled=(device == "cuda"))
    best = 0.0
    for ep in range(epochs):
        t0 = time.time()
        tr = run_epoch(model, train_dl, device, opt, sched, scaler)
        va = run_epoch(model, val_dl, device)
        print(f"ep {ep} train {tr} val {va} {time.time()-t0:.0f}s")
        if va[1] > best:
            best = va[1]; torch.save(model.state_dict(), "best.pt")
    return best

device = "cuda" if torch.cuda.is_available() else "cpu"
torch.manual_seed(SEED)  # seed before the new head is initialized
probe_acc = fit(build_model(freeze_backbone=True), epochs=5, device=device)
torch.manual_seed(SEED)
model = build_model(freeze_backbone=False)
ft_acc = fit(model, epochs=12, device=device)

Each call to fit tracks its own best epoch, so the fine-tuning run overwrites the probe's best.pt; make the file name an argument of fit if you want to keep both checkpoints.

Points to notice. model.train(False) switches batch normalization to its running statistics and disables dropout; forgetting it during evaluation gives noisy, pessimistic numbers. Label smoothing (0.1) replaces the one-hot target with 0.9 on the true class and spreads 0.1 over all classes, which reduces overconfidence on a small dataset. Mixed precision (autocast plus GradScaler) roughly halves activation memory on CUDA; the scaler multiplies the loss before the backward pass so small FP16 gradients do not underflow to zero. The scheduler steps once per batch because T_max is measured in batches. Checkpoint selection uses validation accuracy, never test accuracy.

Milestone 4: Evaluation and error analysis

Load best.pt, run the test set once, and compute the confusion matrix $C$, where $C_{ij}$ is the number of test images whose true class is $i$ and predicted class is $j$. For class $k$:

Macro-F1 is the unweighted mean of $F1_k$ over all classes, so a weak breed cannot hide behind strong ones.

Worked example. Suppose three of the colour-pointed cat breeds have 100 test images each and the slice of the confusion matrix (rows = true, columns = predicted) is:

Birman recall is $81/100 = 0.81$. Birman precision is $81/(5+81+12) = 81/98 = 0.827$. Birman F1 is $2 \times 0.827 \times 0.81 / (0.827 + 0.81) = 0.818$. Doing the same for the others gives Siamese precision $92/98 = 0.939$, recall $0.92$, F1 $0.929$; Ragdoll precision $86/104 = 0.827$, recall $0.86$, F1 $0.843$. Accuracy over the slice is $(92+81+86)/300 = 0.863$ and macro-F1 is $(0.929+0.818+0.843)/3 = 0.864$. The story the matrix tells is more useful than the numbers: 27 of the 41 errors are Birman and Ragdoll mistaken for each other. Both are pale, semi-long-haired cats with dark points; Birmans always have white "gloves" on their paws, but so do mitted Ragdolls, and the paws are often out of frame anyway. That is a data finding, not a model bug.

import numpy as np
from sklearn.metrics import confusion_matrix, classification_report

@torch.no_grad()
def predict(model, loader, device):
    model.eval().to(device)
    ys, ps, confs = [], [], []
    for x, y in loader:
        prob = model(x.to(device)).softmax(1).cpu()
        ys.append(y); ps.append(prob.argmax(1)); confs.append(prob.max(1).values)
    return torch.cat(ys).numpy(), torch.cat(ps).numpy(), torch.cat(confs).numpy()

model.load_state_dict(torch.load("best.pt", map_location=device))
y, p, conf = predict(model, test_dl, device)
labels = list(range(len(classes)))
cm = confusion_matrix(y, p, labels=labels)
print(classification_report(y, p, labels=labels, target_names=classes, digits=3))
off = cm.copy(); np.fill_diagonal(off, 0)
for i, j in zip(*np.unravel_index(np.argsort(off, axis=None)[::-1][:8], off.shape)):
    print(f"{classes[i]:>28} -> {classes[j]:<28} {off[i, j]}")
wrong = np.where(y != p)[0]
print("most confident mistakes:", wrong[np.argsort(-conf[wrong])][:20])

For the error analysis, view the 20 most confident mistakes and tag each with a cause: look-alike breeds, label noise (the annotation looks wrong), unusual pose or occlusion, multiple animals, or the animal is small in the frame. Count the tags. Then check calibration: bin predictions by confidence and compare mean confidence with accuracy in each bin. Optionally, run Grad-CAM on a few mistakes to see which pixels drove the decision.

Milestone 5: Export to ONNX and verify

ONNX is an open format for computation graphs, and ONNX Runtime executes it on CPUs and many accelerators without PyTorch installed. Export in eval mode with a dynamic batch dimension, then prove the exported model matches.

import numpy as np, onnxruntime as ort

model.load_state_dict(torch.load("best.pt", map_location="cpu"))
model.eval().cpu()
dummy = torch.randn(1, 3, 224, 224)
torch.onnx.export(model, (dummy,), "pets_resnet50.onnx",
                  input_names=["image"], output_names=["logits"],
                  dynamic_axes={"image": {0: "batch"}, "logits": {0: "batch"}})

sess = ort.InferenceSession("pets_resnet50.onnx", providers=["CPUExecutionProvider"])
x = torch.randn(8, 3, 224, 224)
with torch.no_grad():
    ref = model(x).numpy()
out = sess.run(None, {"image": x.numpy()})[0]
print("max abs diff:", np.abs(ref - out).max())
print("argmax agree:", (ref.argmax(1) == out.argmax(1)).mean())

Recent PyTorch releases also offer (and in newer versions default to) an exporter built on torch.export, which describes variable dimensions with a dynamic_shapes argument; dynamic_axes is the classic form. If your version warns about or rejects it, follow the torch.onnx documentation for your release rather than guessing. A maximum absolute difference around $10^{-4}$ or smaller on the logits is normal floating-point noise; argmax agreement should be 100% on real test images. Ship the preprocessing contract with the model (resize 256, center crop 224, ImageNet mean and standard deviation, RGB channel order) in a small JSON file, because the most common production bug is preprocessing that differs from training. Measure CPU latency at batch size 1 (median and 95th percentile over 100 runs after 10 warm-up runs).

Evaluation criteria and target metrics

The targets below are this course's thresholds, chosen so that a correct pipeline should clear them; they are not published benchmark results. If you are far below them, suspect a bug (wrong normalization, evaluation in train mode, label misalignment) before you suspect the model.

Grading rubric

Stretch goals

• Test-time augmentation (average predictions over the image and its horizontal flip) and measure the gain.
• Quantize the ONNX model to INT8 with ONNX Runtime's quantization tools and report the accuracy and latency trade-off.
• Train a two-level model: cat versus dog first, then breed, and compare error patterns.
• Add an "unknown" decision: reject images whose top probability is below a threshold chosen on validation, and report coverage versus accuracy.
• Fine-tune ViT-B/16 or a ConvNeXt and compare accuracy per parameter and per millisecond.

Write-up template

1. Problem and users. Who uses the prediction and what a mistake costs.
2. Data. Source, licence, split sizes, how splits were made, class balance.
3. Method. Architecture, pretrained weights, augmentation, optimizer, learning rates, schedule, epochs, hardware, wall-clock time, library versions.
4. Results. Table of probe versus fine-tune on validation and test; learning curves.
5. Error analysis. Top confused pairs, tagged error causes with counts, calibration.
6. Deployment. ONNX file size, equivalence check, latency, preprocessing contract.
7. Limitations and next steps. What the model should not be used for (for example, mixed breeds or species outside the 37 classes).

Common pitfalls

• Evaluating with augmentation or in train mode, which makes validation noisy and pessimistic.
• Choosing hyperparameters or checkpoints on the test set, which makes the reported number optimistic.
• Using one learning rate for everything: too high wrecks pretrained features, too low leaves the new head undertrained.
• Normalizing with the wrong statistics or feeding BGR images at inference.
• Reporting only accuracy when the business cost is concentrated in a few classes.

In practice

Teams start with the strongest pretrained backbone their latency budget allows, run a linear probe as a smoke test, and fine-tune with conservative learning rates. Most of the improvement after that comes from data work driven by error analysis: fixing labels, collecting more examples of confused classes, and agreeing on what the model should do with out-of-scope images. Exported models are versioned together with their preprocessing and an equivalence test so that a change on either side is caught before release.