Project: train and deploy an image classifier

The brief

Build a classifier for 5–20 visual categories you care about (plant species in your area, kinds of food, tools in a workshop, hand-written symbols, dog breeds from a public dataset). Target: a model that beats a reasonable baseline, an error analysis you can explain, and a live demo. Budget: 6–10 hours over a week on a free Colab/Kaggle GPU. If collecting images is hard, use Oxford-IIIT Pets (37 breeds, 7,400 images) or a 10-class subset of Food-101.

Milestone 1 — Data (1–2 h)

Collect 100–500 images per class (your own photos, a public dataset, or a careful web scrape you have the right to use). Put them in data/train/<class>/ folders so torchvision.datasets.ImageFolder works. Then:

• Deduplicate with perceptual hashes (imagehash) — near-duplicates across train and validation inflate accuracy, the most common beginner leak (3.5 Evaluating models: metrics, validation and the bias–variance trade-off).
• Split 70/15/15 at the level of the source (same photo session or same object stays in one split), stratified by class.
• Look at the data: a grid of 64 random images per class; note label noise and class imbalance now, not after training.

Milestone 2 — Baseline and model (2–3 h)

Start with a trivial baseline (majority class; then logistic regression on frozen convnext_tiny embeddings — often strong, 5.4 Transfer learning and training CNNs in practice). Then fine-tune:

import timm, torch, torch.nn as nn
from torchvision import datasets
model = timm.create_model("convnext_tiny", pretrained=True, num_classes=NUM_CLASSES).cuda()
cfg = timm.data.resolve_data_config({}, model=model)
train_ds = datasets.ImageFolder("data/train", timm.data.create_transform(**cfg, is_training=True, auto_augment="rand-m9-mstd0.5"))
val_ds   = datasets.ImageFolder("data/val",   timm.data.create_transform(**cfg, is_training=False))
train_dl = torch.utils.data.DataLoader(train_ds, batch_size=32, shuffle=True, num_workers=4)
val_dl   = torch.utils.data.DataLoader(val_ds, batch_size=64, num_workers=4)
opt = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=0.05)
sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=EPOCHS * len(train_dl))
loss_fn = nn.CrossEntropyLoss(label_smoothing=0.1)
for epoch in range(EPOCHS):
    model.train()
    for x, y in train_dl:
        with torch.autocast("cuda", dtype=torch.bfloat16):
            loss = loss_fn(model(x.cuda()), y.cuda())
        opt.zero_grad(set_to_none=True); loss.backward(); opt.step(); sched.step()
    model.eval(); correct = 0
    with torch.no_grad():
        for x, y in val_dl:
            correct += (model(x.cuda()).argmax(1).cpu() == y).sum().item()
    print(epoch, loss.item(), correct / len(val_ds))
torch.save(model.state_dict(), "model.pt")

Log every run (a CSV or Weights & Biases, 2.5 Tooling: Git, GPUs, experiments and the Hugging Face Hub): backbone, resolution, LR, augmentation, epochs, validation accuracy. Change one thing at a time. Expect fine-tuning to beat the frozen-embedding baseline by a few points when you have a few hundred images per class, and to be close when you have fewer.

Milestone 3 — Error analysis (1–2 h)

• Confusion matrix on the validation set; name the top three confused pairs and look at those images.
• Top losses: the 32 highest-loss validation images — fix label errors, remove junk, note genuinely hard cases.
• Grad-CAM (5.9 Interpretability, robustness and adversarial examples) on 10 correct and 10 wrong predictions: is the model looking at the object or at a shortcut (background, watermark, a hand)?
• Calibration: reliability diagram; fit a temperature so your demo's confidence means something.
• Shift test: 20 photos taken under new conditions. Report the drop honestly. Only now touch the test set, once, for the final number.

Milestone 4 — Deploy (1–2 h)

Export the weights (and the class names and preprocessing config!) and write app.py:

import gradio as gr, timm, torch, json
from PIL import Image
classes = json.load(open("classes.json"))
model = timm.create_model("convnext_tiny", pretrained=False, num_classes=len(classes))
model.load_state_dict(torch.load("model.pt", map_location="cpu")); model.eval()
cfg = timm.data.resolve_data_config({}, model=model)
tf = timm.data.create_transform(**cfg, is_training=False)
T = 1.8   # temperature from calibration

def predict(img: Image.Image):
    with torch.no_grad():
        probs = torch.softmax(model(tf(img).unsqueeze(0)) / T, dim=1)[0]
    return {c: float(p) for c, p in zip(classes, probs)}

gr.Interface(predict, gr.Image(type="pil"), gr.Label(num_top_classes=3),
             title="My classifier", examples=["examples/a.jpg", "examples/b.jpg"]).launch()

Create a Space on Hugging Face (SDK: Gradio), push app.py, model.pt, classes.json, a requirements.txt (gradio, timm, torch) and a few example images. CPU Spaces are free and fine for a ConvNeXt-Tiny at ~100 ms per image. Stretch goals: export to ONNX for faster CPU inference, add a Grad-CAM overlay output, or ship a TorchScript model inside a small FastAPI service.

Milestone 5 — Write-up (1 h)

A README with: the problem and data (sizes, splits, how leakage was prevented), the baseline and final results in a table (with the test set used once), the three most informative failure cases with Grad-CAM, the shift-test result, limitations, and the demo link. Honest beats impressive: an 88% model with a clear error analysis is a better portfolio piece than a 97% model with hidden leakage.

Rubric

Where this shows up later

• The same fine-tune → evaluate → deploy loop for an LLM — 8.13 Project: fine-tune a small LLM with LoRA and ship a chat app, a RAG app — 9.11 Project: a cited RAG assistant over your own documents, with evals and an agent — 10.12 Project: build a research agent with tools, an MCP server and an eval harness.
• Portfolio and write-up standards — 11.3 Portfolio project ladder; experiment tracking — 2.5 Tooling: Git, GPUs, experiments and the Hugging Face Hub; evaluation discipline — 3.5 Evaluating models: metrics, validation and the bias–variance trade-off.