Image Classification Pipeline

An image classification pipeline converts image files or image-like tensors into class probabilities and decisions. In this PyTorch course, the pipeline is the bridge between convolutional model architecture and evidence that the model is solving the right visual task. The outcome is not just a trained network. The outcome is a repeatable path that loads images, applies identical label rules, transforms pixels into tensors, batches examples, computes logits, updates parameters during training, and reports honest validation metrics.

Purpose and Outcome

Classification answers one question per image: which label from a fixed set best describes the image? A pipeline makes that question executable. For a cat-versus-dog classifier, the pipeline defines where cat and dog examples come from, how filenames become numeric labels, whether pixels are resized or normalized, what tensor shape the model receives, how logits map to probabilities, and how accuracy or confusion matrices are computed. By the end of this lesson you should be able to inspect every stage of that path, build a small version in PyTorch, and diagnose common mismatches before they become misleading scores.

How the Pipeline Works Internally

PyTorch image models normally consume batches shaped [N, C, H, W]: batch size, channels, height, and width. Many image files are stored as height by width by channel, and pixel values often arrive as unsigned integers from 0 to 255. The pipeline must reorder axes, convert type, and scale values before the model sees them. A normalization step then subtracts per-channel means and divides by per-channel standard deviations so the optimizer receives inputs with a more stable numeric range.

The Dataset owns indexing. Its __len__ tells PyTorch how many examples exist, and its __getitem__ returns one image-label pair. The DataLoader owns batching, shuffling, worker processes, and optional pinned memory for faster host-to-device transfer. The model receives a batch and emits logits, not probabilities. Logits are unnormalized scores with shape [N, num_classes]. CrossEntropyLoss combines log_softmax and negative log likelihood, so training code should pass raw logits and integer class indices, not one-hot vectors or already-softmaxed probabilities.

Training mode and evaluation mode change behavior for layers such as dropout and batch normalization. During training, the loop enables gradients, computes loss, calls backward(), and steps the optimizer. During validation or inference, the loop uses model.eval() and torch.no_grad() so layers behave consistently and PyTorch avoids storing the autograd graph. This distinction is part of the pipeline, not a cosmetic optimization.

API Anatomy

Example 1: Tensor Shape and Pixel Scale

The first example starts where many bugs start: pixels. A single RGB image with 4 by 4 pixels is represented as [3, 4, 4]. Adding a batch dimension makes it acceptable to a convolutional model. The expected maximum value is 0.184 because the largest stored pixel is 47 and the code divides by 255.

import torch

image_uint8 = torch.arange(3 * 4 * 4, dtype=torch.uint8).reshape(3, 4, 4)
image_float = image_uint8.float() / 255.0
batch = image_float.unsqueeze(0)

print(batch.shape)
print(round(float(batch.max()), 3))

This tiny inspection catches two high-value mistakes: forgetting the batch dimension and feeding integer pixels directly into a model that was trained on floats. Real pipelines should assert the shape and dtype near the data boundary, because a later convolution error usually hides the original cause.

Example 2: Dataset and DataLoader Contract

A dataset returns individual examples; a loader combines them. The labels here are deterministic because shuffling is disabled. The first batch contains two images with shape (2, 3, 8, 8) and labels [0, 1].

import torch
from torch.utils.data import Dataset, DataLoader

class TinyImageSet(Dataset):
    def __init__(self) -> None:
        self.images = torch.stack([
            torch.zeros(3, 8, 8),
            torch.ones(3, 8, 8),
            torch.full((3, 8, 8), 0.5),
        ])
        self.labels = torch.tensor([0, 1, 1])

    def __len__(self) -> int:
        return len(self.labels)

    def __getitem__(self, index: int) -> tuple[torch.Tensor, torch.Tensor]:
        return self.images[index], self.labels[index]

loader = DataLoader(TinyImageSet(), batch_size=2, shuffle=False)
images, labels = next(iter(loader))
print(tuple(images.shape))
print(labels.tolist())

The contract matters because transforms usually run inside __getitem__. Training transforms may include random crops, flips, or color jitter, while validation transforms should be deterministic. Mixing those policies is a common way to make validation unstable or to overstate model quality.

Example 3: Logits, Probabilities, and Loss

The next step is model output. The network emits three logits per image because it has three classes. Softmax converts each row to probabilities whose sum is 1.0, but the loss receives the logits directly. The printed scalar rank 0 shows that cross entropy returns one scalar for the batch by default.

import torch
import torch.nn as nn
import torch.nn.functional as F

torch.manual_seed(7)
model = nn.Sequential(
    nn.Conv2d(3, 4, kernel_size=3, padding=1),
    nn.ReLU(),
    nn.AdaptiveAvgPool2d((1, 1)),
    nn.Flatten(),
    nn.Linear(4, 3),
)

x = torch.randn(2, 3, 8, 8)
y = torch.tensor([0, 2])
logits = model(x)
loss = F.cross_entropy(logits, y)
probability_sums = torch.softmax(logits, dim=1).sum(dim=1)

print(tuple(logits.shape))
print([round(float(v), 4) for v in probability_sums])
print(loss.ndim)

Do not put softmax in the model head when using F.cross_entropy. Doing so changes the numeric surface that the loss expects and can slow or destabilize learning. Apply softmax only when reporting probabilities or thresholding downstream decisions.

Example 4: A Minimal Train-Evaluate Loop

This synthetic grayscale task is deliberately simple: black images are class 0 and white images are class 1. A linear classifier over flattened pixels can learn it. After several optimizer steps, the expected predictions are [0, 0, 1, 1].

import torch
import torch.nn as nn
import torch.nn.functional as F

torch.manual_seed(0)
images = torch.zeros(4, 1, 4, 4)
images[2:] = 1.0
labels = torch.tensor([0, 0, 1, 1])
model = nn.Sequential(nn.Flatten(), nn.Linear(16, 2))
optimizer = torch.optim.SGD(model.parameters(), lr=0.5)

for _ in range(25):
    optimizer.zero_grad()
    logits = model(images)
    loss = F.cross_entropy(logits, labels)
    loss.backward()
    optimizer.step()

with torch.no_grad():
    predictions = model(images).argmax(dim=1)

print(predictions.tolist())

The important pieces are the order of operations. Clear gradients before the forward pass, compute loss, backpropagate, then update parameters. For evaluation, switch off gradient tracking and avoid optimizer calls. The same structure scales to convolutional networks and real datasets.

Design Choices and Trade-offs

Image size is the first major trade-off. Larger images preserve detail but increase memory, latency, and overfitting risk. Smaller images train faster but may erase the signal, such as fine cracks in inspection data or small lesions in medical images. Cropping policies also matter: center crops are deterministic but may remove off-center objects, while random resized crops improve robustness but can make rare features harder to learn.

Class design is another choice. Mutually exclusive classes fit standard cross entropy. Multi-label tasks, where one image may contain several tags, require sigmoid outputs and a binary loss instead. Class imbalance changes the evaluation story: high accuracy can hide a model that ignores rare classes. Use per-class precision, recall, confusion matrices, or weighted sampling when the label distribution is uneven.

Augmentation should match the domain. Horizontal flips are often acceptable for natural images, but they can be wrong for text, road signs, or medical images where orientation carries meaning. Normalization should match the data used to train the model. If using a pretrained backbone, use the same channel convention and normalization statistics expected by that backbone, otherwise early feature activations shift away from the distribution the filters learned.

Failure Modes and Troubleshooting

• Symptom: the first convolution reports a channel mismatch. Cause: tensors are shaped [N, H, W, C] or grayscale images are missing a channel dimension. Diagnose: print one batch shape before the model. Correct: convert to [N, C, H, W] and explicitly handle RGB versus grayscale.
• Symptom: training loss decreases but validation accuracy is near random. Cause: label mapping differs between train and validation folders, or validation transforms are destructive. Diagnose: print class_to_idx, inspect a few image-label pairs, and run a tiny overfit test on one batch. Correct: share one class mapping and make validation transforms deterministic.
• Symptom: loss becomes nan. Cause: too high a learning rate, invalid input values, or exploding activations. Diagnose: check torch.isfinite on inputs, logits, and loss; lower the learning rate for one run. Correct: fix corrupt data, reduce the learning rate, normalize inputs, or add gradient clipping for deep models.
• Symptom: GPU usage is low while CPU is busy. Cause: the loader cannot prepare batches fast enough. Diagnose: time data loading separately from forward and backward passes. Correct: tune num_workers, batch size, storage layout, and pinned memory where supported.

Security, Performance, and Reliability

Image inputs are untrusted data in many systems. Limit accepted formats and image dimensions before decoding large files, because malformed or huge images can exhaust memory. Keep dataset paths explicit so a training job does not silently include private or mislabeled folders. Store class names and preprocessing configuration with the model artifact; otherwise an exported model can make numerically valid but semantically wrong predictions.

Performance depends on both compute and input throughput. Batch size should be large enough to use the accelerator but small enough to avoid out-of-memory failures. Mixed precision can improve throughput on supported hardware, but verify accuracy because small or imbalanced datasets may react badly to numeric changes. Reliability comes from reproducible splits, saved seeds when feasible, versioned data manifests, and validation metrics computed with the model in evaluation mode.

Hands-on Lab

Prerequisites: Python with PyTorch installed, a shell, and a writable scratch directory outside the published lesson source. A GPU is optional. Use the examples above as the starting script.

1. Create a small script named pipeline_lab.py in your scratch directory and paste Example 4.
2. Run it and verify that the predictions print [0, 0, 1, 1].
3. Add a validation block that computes accuracy as (predictions == labels).float().mean() and print it. The expected value is 1.0.
4. Change the labels to [0, 1, 0, 1] and rerun. The simple intensity rule is now contradictory, so accuracy should no longer be a trustworthy sign of learning.
5. Restore the original labels, then insert print(images.shape, images.dtype, labels.shape) before training to verify the data contract.

Cleanup: remove the scratch script and any generated cache files. No model artifact needs to be retained for this lab. If you adapted the lab to real images, also remove temporary resized copies unless they are part of a documented dataset version.

Assessment Exercises

1. You inherit a classifier with 98 percent accuracy but poor results on one class. Which metrics and data slices would you inspect, and what pipeline bug could create this pattern?
2. A model trained with random crops performs unpredictably during validation. Explain where transform policy belongs and how you would separate training from validation behavior.
3. Why does CrossEntropyLoss expect logits and integer labels? Describe the mistake caused by applying softmax before the loss.
4. Design a quick overfit test for one image batch. What result would convince you that the model, labels, and optimizer are wired correctly?
5. For a pretrained RGB backbone, what should you check before feeding grayscale or differently normalized images?

Summary

A PyTorch image classification pipeline is a sequence of enforceable contracts: decode images, transform them into normalized [N, C, H, W] tensors, batch them, compute logits, optimize with the right loss, and evaluate with deterministic preprocessing. Most hard failures are shape, dtype, label, mode, or throughput problems. Treat each stage as inspectable, keep training and validation policies separate, and verify the smallest complete loop before scaling to a larger dataset or model.