Module 8 · Deployment and Capstone

Capstone: Train and Deploy an Image Classifier

Lesson 32 of 32 11 min read Updated

This capstone ties together the course by turning a PyTorch image classifier into a repeatable artifact that can answer real inference requests. The outcome is not only a trained neural network. It is a folder of versioned assets, a preprocessing contract, a label mapping, an evaluation record, and a small service interface that converts an uploaded image into class probabilities.

The central idea is simple: during training, the model learns parameters that map normalized image tensors to logits. During deployment, the same tensor construction must be recreated without the training-only pieces such as shuffling, random crops, gradient tracking, and optimizer state. Most classifier deployment failures come from breaking that continuity. A model trained on resized RGB images normalized with one mean and standard deviation will behave unpredictably if production sends BGR images, a different crop, or class indices from a different directory order.

Classifier Pipeline Anatomy

An image classification system has five coupled parts. The dataset defines examples and class names. The transform converts files into tensors with shape [channels, height, width]. The model maps batches shaped [batch, channels, height, width] to logits shaped [batch, classes]. The loss function compares logits with integer class ids. The deployment wrapper reproduces inference preprocessing, applies softmax, and translates numeric indices back to human-readable labels.

PyTorch does not automatically remember every training assumption. A state_dict stores tensor values for named parameters and buffers, not the Python class definition, preprocessing pipeline, or semantic label names. That is why a useful classifier artifact usually contains more than model.pt. Store the architecture name or module code revision, input size, normalization constants, label mapping, metric summary, and creation time. This metadata is what lets another process reconstruct the same computation later.

Mechanics Inside Training and Inference

For supervised classification, each training step starts with a mini-batch of images and target ids. The model produces logits, which are unnormalized class scores. CrossEntropyLoss combines log_softmax and negative log likelihood, so you pass raw logits to it. After loss.backward(), PyTorch autograd has populated gradients for parameters involved in the forward pass. The optimizer reads those gradients and updates weights. Clearing gradients before the next batch avoids accidentally accumulating gradients across independent batches.

Inference uses the same forward computation but a different execution context. model.eval() changes modules such as dropout and batch normalization to evaluation behavior. torch.inference_mode() disables autograd bookkeeping, reducing memory use and preventing accidental parameter updates. The service should receive bytes, decode an image, convert it to RGB, resize and normalize it exactly as expected, add a batch dimension, run the model, and return labels and probabilities.

Configuration Choices

The simplest capstone can train a small convolutional network from scratch. A more realistic project often fine-tunes a pretrained backbone because early visual filters transfer well across many natural image tasks. Fine-tuning trades extra dependency on a known architecture for faster convergence and better accuracy on small datasets. Freezing the backbone is faster and reduces overfitting risk, but it may underperform when the new images differ strongly from the pretraining domain. Unfreezing more layers gives the optimizer more capacity, but it also needs more data, lower learning rates, and closer validation monitoring.

Class mapping is another important design choice. Directory-based datasets commonly infer labels from folder names in sorted order. If cat was id 0 during training, production must not later treat id 0 as dog. Save idx_to_class with the model. Do not regenerate it casually from a changed filesystem.

Example 1: Save the Label Contract

This first example shows why label metadata belongs beside the model. It builds a deterministic mapping, serializes it, and verifies that index lookup is stable. Expected output is class 1: daisy.

import json
from pathlib import Path

classes = ["dandelion", "daisy", "rose"]
class_to_idx = {name: index for index, name in enumerate(classes)}
idx_to_class = {str(index): name for name, index in class_to_idx.items()}

artifact_dir = Path("classifier_artifact")
artifact_dir.mkdir(exist_ok=True)
metadata = {
    "input_size": 224,
    "mean": [0.485, 0.456, 0.406],
    "std": [0.229, 0.224, 0.225],
    "idx_to_class": idx_to_class,
}
metadata_path = artifact_dir / "metadata.json"
metadata_path.write_text(json.dumps(metadata, indent=2), encoding="utf-8")

loaded = json.loads(metadata_path.read_text(encoding="utf-8"))
print(f"class 1: {loaded['idx_to_class']['1']}")

The mechanism is deliberately plain. The keys in JSON become strings, so the deployment code should either use string keys or convert them back to integers. This example also records image size and normalization constants because those values are part of the inference contract.

Example 2: One Training Step

The next example uses random tensors to demonstrate the training loop anatomy without requiring image files. It defines a tiny convolutional classifier, runs one optimization step, and prints the output shape. Expected output is logits shape: (4, 3).

import torch
from torch import nn

class TinyImageClassifier(nn.Module):
    def __init__(self, num_classes: int) -> None:
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 8, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d((1, 1)),
        )
        self.classifier = nn.Linear(8, num_classes)

    def forward(self, images: torch.Tensor) -> torch.Tensor:
        features = self.features(images)
        flattened = features.flatten(1)
        return self.classifier(flattened)

model = TinyImageClassifier(num_classes=3)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()

images = torch.randn(4, 3, 64, 64)
targets = torch.tensor([0, 1, 2, 1])

model.train()
optimizer.zero_grad(set_to_none=True)
logits = model(images)
loss = criterion(logits, targets)
loss.backward()
optimizer.step()

print(f"logits shape: {tuple(logits.shape)}")

The shape tells you the model is producing one row per image and one column per class. The targets are integer class ids, not one-hot vectors, because CrossEntropyLoss expects class indices. In a real project, replace the random tensors with a DataLoader fed by training transforms that include augmentation, while validation transforms stay deterministic.

Example 3: Deployable Prediction Function

The third example isolates the core of an API handler: convert logits to probabilities and return the top class using the saved mapping. Expected output is {'label': 'rose', 'probability': 0.936}.

import torch

def predict_from_logits(logits: torch.Tensor, idx_to_class: dict[str, str]) -> dict[str, object]:
    if logits.ndim != 2 or logits.shape[0] != 1:
        raise ValueError("expected logits with shape [1, classes]")
    probabilities = torch.softmax(logits, dim=1)
    probability, index = torch.max(probabilities, dim=1)
    label = idx_to_class[str(index.item())]
    return {"label": label, "probability": round(float(probability.item()), 3)}

logits = torch.tensor([[0.2, -1.0, 3.1]])
labels = {"0": "dandelion", "1": "daisy", "2": "rose"}
print(predict_from_logits(logits, labels))

Production code would place image decoding and model execution before this function, but the output step is the same. It is usually better to return the top few classes rather than only one. A low top probability or a narrow margin between the first and second class can signal that the request should be reviewed instead of accepted automatically.

Failure Modes and Troubleshooting

Symptom: validation accuracy is high, but deployed predictions are poor. Cause: preprocessing drift, such as missing RGB conversion or using a different resize. Diagnostic steps: save one validation image tensor during training, run the deployment preprocessing on the same source image, and compare shape, dtype, min, max, channel order, and normalized values. Correction: move preprocessing constants into metadata and write a test that compares production preprocessing with the training transform for a fixed image.

Symptom: the API returns plausible probabilities with the wrong class names. Cause: the class directory order changed or labels were rebuilt during deployment. Diagnostic steps: inspect the saved idx_to_class, compare it with the training run record, and test a known image whose class is unambiguous. Correction: load labels only from the artifact metadata and reject startup if the number of model outputs does not match the number of saved labels.

Symptom: GPU memory grows during serving. Cause: inference runs with gradient tracking or stores tensors from requests. Diagnostic steps: check for missing torch.inference_mode(), retained outputs in global lists, and request handlers that call backward(). Correction: wrap prediction in inference mode, detach values before logging, and avoid retaining request tensors.

Symptom: training loss falls while validation loss rises. Cause: overfitting, label leakage, or an overly aggressive fine-tuning setup. Diagnostic steps: verify split construction, inspect near-duplicate images across train and validation sets, plot per-class metrics, and compare frozen versus unfrozen backbones. Correction: clean the split, add augmentation, reduce trainable parameters, or stop from the best validation checkpoint.

Security, Performance, and Reliability

Image endpoints need input limits. Reject files that are too large, enforce supported formats, and handle decode failures without crashing the process. Do not trust client-provided filenames or content types. If images are stored for audit or debugging, define retention and avoid logging raw personal data unnecessarily.

Performance depends on batching, device placement, and model size. A small CPU model may be more reliable and cheaper than a large GPU model for low traffic. For high traffic, keep the model loaded at process startup, avoid reloading weights per request, and measure p95 latency with representative image sizes. Reliability improves when artifacts are immutable: deploy a new versioned directory, health-check it, and roll back by pointing the service to the previous artifact.

Hands-On Lab

Prerequisites: Python with PyTorch installed, a small image dataset arranged as data/train/class_name/*.jpg and data/val/class_name/*.jpg, and enough disk space to save one artifact directory. The dataset can be flowers, products, defects, or any small set of visually distinct classes.

  1. Create a reproducible split. Keep validation images out of training augmentation decisions and record the class names.
  2. Build transforms. Training can use random crop or flip; validation and serving should use deterministic resize, center crop if needed, tensor conversion, and normalization.
  3. Choose the model. Start with a pretrained backbone when the dataset is small. Replace the final classifier layer so the output width equals the number of classes.
  4. Train for a small number of epochs. Track training loss, validation loss, accuracy, and per-class accuracy. Save the checkpoint with the best validation metric, not necessarily the last epoch.
  5. Create an artifact directory containing model weights, metadata JSON, and a short metrics file. Include idx_to_class, input size, normalization constants, architecture name, and validation score.
  6. Write a prediction function that loads the artifact once, decodes one image, runs model.eval() with torch.inference_mode(), and returns top probabilities.
  7. Wrap the function in a small API or command-line script. Validate file size and image decoding before inference.

Verification: run prediction on three validation images whose labels you know, confirm the returned labels match the saved mapping, and compare one API result with a direct notebook prediction from the same artifact. Check that a text file renamed as .jpg is rejected cleanly. Measure average and p95 latency over at least twenty local requests.

Cleanup or rollback: delete temporary uploaded images, remove failed artifact directories, and keep the previous known-good artifact available. If a new model performs worse, switch the service configuration back to the previous artifact path and restart or reload the process.

Assessment Exercises

  1. You trained with random resized crops but serve with a simple stretch resize. What symptoms would you expect, and how would you prove preprocessing drift is the cause?
  2. A checkpoint has four output logits, but metadata.json contains five class names. Should the service start? Explain the invariant being protected.
  3. Your validation accuracy is 95 percent, but one class has 40 percent recall. What deployment decision would you make, and what extra metric or workflow would you add?
  4. Compare freezing a pretrained backbone with fine-tuning all layers for a dataset of 300 images. What risks and benefits matter most?
  5. Design a rollback test for an image classifier API. What request would you send before and after rollback, and what must remain stable?

Summary

A deployable PyTorch image classifier is a coordinated system, not a lone weight file. Training learns the mapping from image tensors to logits; deployment must reconstruct the same preprocessing, architecture, weights, and label semantics. Save metadata with the model, validate shapes and class counts at startup, use evaluation mode and inference mode for serving, and test both known images and malformed inputs. That discipline turns the capstone from a notebook result into a classifier that can be inspected, repeated, and safely replaced.

Progress is saved in this browser.