Image Classification — Module 7.2

How This Lesson Fits the Module

Module 7.1 built the convolutional toolkit: convolution, filters, pooling, CNNs, and transfer learning. Module 7.2 now applies that toolkit to the three canonical vision tasks.

Image classification is the foundational task and the entry point for the module: given a whole image, predict one label (or a ranked set of labels). It defines the vocabulary—logits, softmax, top-k accuracy, backbones—that the harder tasks (object detection and image segmentation) build on. Master classification first, and the rest of computer vision becomes an exercise in "classify where" and "classify every pixel."

Learning Objectives

By the end of this lesson, students should be able to:

Define image classification and distinguish it from detection and segmentation.
Explain the classification pipeline: backbone → global pooling → classifier head → softmax.
Differentiate single-label, multi-class, and multi-label classification and choose the right loss.
Load and fine-tune a pretrained torchvision model (ResNet) for a custom dataset.
Compute and interpret top-1 / top-5 accuracy, precision, recall, and the confusion matrix.
Diagnose common failure modes: class imbalance, data leakage, and wrong-loss bugs.

What Is Image Classification?

Image classification maps an entire image to a category. The model never localizes anything—it answers a single question: "What is this a picture of?" A photo of a dog on a beach is simply dog (or Labrador in a fine-grained setting). This is the simplest of the three vision tasks precisely because the output is a single vector of class scores rather than boxes or per-pixel maps.

Definition — Image Classification

Image classification is the task of assigning one or more predefined category labels to an entire input image. A classifier f(x) produces a vector of logits over C classes; a normalization (softmax or sigmoid) turns logits into probabilities, and the arg-max (or a threshold) yields the prediction.

The Three Vision Tasks at a Glance

Before diving in, it helps to see where classification sits relative to its siblings in this module.

Task	Question Answered	Output	Typical Loss	Metric
Classification	What is in the image?	1 label per image	Cross-entropy	Top-1 / Top-5 accuracy
Detection	What and where (boxes)?	Boxes + labels	Cls + box regression	mAP @ IoU
Segmentation	Which pixels belong to what?	Per-pixel mask	Pixel cross-entropy / Dice	mIoU / Dice

The Classification Pipeline

Nearly every modern image classifier shares the same anatomy. A convolutional (or transformer) backbone turns pixels into a stack of feature maps; a global pooling layer collapses spatial dimensions into a single feature vector; and a small classifier head (usually one linear layer) projects that vector onto C class logits.

1. Input

Normalized image tensor, shape (B, 3, H, W).

2. Backbone

Conv stack extracts features, e.g. (B, 512, 7, 7).

3. Global Pool

AdaptiveAvgPool collapses to (B, 512).

4. Head

Linear layer → logits (B, C).

5. Softmax

Logits → class probabilities.

Definition — Backbone

A backbone is the feature-extraction body of a network (e.g. ResNet-50 minus its final classifier). The same backbone can be reused across tasks: classification adds a linear head, detection adds box heads, and segmentation adds a decoder. This reuse is why transfer learning is so powerful.

Types of Classification

The word "classification" hides three distinct problem shapes, and choosing the wrong one silently breaks training.

Binary

Two classes (cat vs. not-cat)
1 logit + sigmoid
BCEWithLogitsLoss

Multi-class

One label from C classes
C logits + softmax
CrossEntropyLoss

Multi-label

Several labels can co-occur
C logits + sigmoid (independent)
BCEWithLogitsLoss

Critical Mistake — Softmax on Multi-label

Softmax forces probabilities to sum to 1, encoding "exactly one class is correct." For a multi-label problem (an image tagged both beach and sunset) this is wrong—use per-class sigmoid + BCEWithLogitsLoss instead. Using softmax here makes the labels compete and caps recall.

Fine-Tuning a Pretrained Classifier in PyTorch

In practice you rarely train a classifier from scratch. You take a backbone pretrained on ImageNet and fine-tune it—the transfer learning workflow from Module 7.1. Here we adapt a ResNet-18 to a custom 5-class dataset.

import torch
import torch.nn as nn
from torchvision import models, transforms
from torchvision.models import ResNet18_Weights

# 1. Load a backbone pretrained on ImageNet
weights = ResNet18_Weights.DEFAULT
model = models.resnet18(weights=weights)

# 2. Replace the 1000-class head with our own (5 classes)
num_classes = 5
model.fc = nn.Linear(model.fc.in_features, num_classes)

# 3. Use the weights' own preprocessing transforms
preprocess = weights.transforms()   # resize, center-crop, normalize

# 4. Standard multi-class setup
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)

# 5. One training step
def train_step(images, labels):
    optimizer.zero_grad()
    logits = model(images)              # (B, 5) raw scores
    loss = criterion(logits, labels)   # labels are class indices, not one-hot
    loss.backward()
    optimizer.step()
    return loss.item()

Critical Mistake — Applying Softmax Before CrossEntropyLoss

nn.CrossEntropyLoss expects raw logits and applies log_softmax internally. Passing already-softmaxed probabilities double-normalizes the scores, shrinks gradients, and slows or breaks learning. Feed logits directly.

Inference and Reading the Output

model.eval()
with torch.no_grad():
    logits = model(image_batch)             # (B, C)
    probs = torch.softmax(logits, dim=1)    # probabilities per class
    top5 = probs.topk(5, dim=1)             # top-5 values + indices
    pred = probs.argmax(dim=1)              # top-1 predicted class

# Interpretability tip: high top-1 prob != correct.
# Always inspect the confusion matrix, not just accuracy.
print(pred, top5.indices)

Evaluating a Classifier

Accuracy alone can be deeply misleading. On a dataset that is 95% "healthy" scans, a model that always predicts "healthy" scores 95% accuracy while catching zero disease. Report the metrics below together.

Metric	Definition	Best For
Top-1 Accuracy	Fraction where arg-max = truth	Balanced, single-label
Top-5 Accuracy	Truth in top-5 predictions	Many fine-grained classes (ImageNet)
Precision	TP / (TP + FP)	Cost of false alarms is high
Recall	TP / (TP + FN)	Missing a positive is costly
F1 Score	Harmonic mean of P and R	Imbalanced classes
Confusion Matrix	Per-class predicted vs. actual	Diagnosing which classes confuse

Strengths and Tradeoffs

Why classification is the starting point

Simplest output space: one vector of scores.
Cheap labels—just one tag per image.
Pretrained backbones transfer to every other vision task.
Fast to train and easy to benchmark.

Where it falls short

No localization—cannot say where the object is.
Assumes one dominant subject per image.
Struggles with cluttered scenes containing many objects.
Sensitive to class imbalance and spurious correlations.

Common Misconception: “A confident (high-probability) prediction means the model is correct.”

Reality: Softmax outputs are not calibrated probabilities. Deep networks are frequently over-confident, assigning 0.99 to wrong answers on out-of-distribution inputs. Calibration (temperature scaling) and the confusion matrix matter more than raw confidence.

Common Misconception: “Classification and detection are basically the same thing.”

Reality: Classification returns a single label for the whole image; detection returns a variable number of boxes, each with its own label and confidence. The output structure, loss, and metrics are all different.

Module 6 Synthesis Training this classifier reuses everything from Vol 06: cross-entropy loss, AdamW, the training loop, and batch normalization inside the backbone.

Knowledge Check

Short Answer: What single object does an image classifier output for one image? Answer: A vector of class scores (logits) leading to one predicted label.
True/False: nn.CrossEntropyLoss expects raw logits, not softmaxed probabilities. Answer: True.
Multiple Choice: For a multi-label problem you should use: (a) softmax + CrossEntropy, (b) sigmoid + BCEWithLogitsLoss, (c) argmax only, (d) MSE. Answer: (b).
Short Answer: Name the three parts of the classification pipeline after the input. Answer: Backbone, global pooling, classifier head (then softmax).
True/False: Top-5 accuracy is stricter than top-1 accuracy. Answer: False—top-5 is more lenient.
Multiple Choice: On a 95%-negative dataset, which metric best exposes a lazy "always negative" model? (a) accuracy, (b) recall, (c) parameter count, (d) learning rate. Answer: (b) recall.
Short Answer: Why replace model.fc when fine-tuning a torchvision ResNet? Answer: The pretrained head outputs 1000 ImageNet classes; we need a head sized to our class count.
True/False: Softmax outputs are always well-calibrated probabilities. Answer: False—deep nets are often over-confident.
Multiple Choice: Which task adds bounding boxes on top of classification? (a) Segmentation, (b) Detection, (c) Pooling, (d) Dropout. Answer: (b).
Short Answer: Which Module 7.3 model families are classic classification backbones? Answer: ResNet, VGG, EfficientNet, ViT.

Key Takeaways

Classification assigns one (or several) labels to a whole image—no localization.
The pipeline is backbone → global pool → linear head → softmax/sigmoid.
Choose the loss by problem shape: CrossEntropy (multi-class) vs. BCEWithLogits (multi-label).
Fine-tune pretrained torchvision backbones instead of training from scratch.
Report precision, recall, F1, and a confusion matrix—never accuracy alone.
Next: Object Detection extends classification with "where."

Trainer’s Guide

Live demo: Fine-tune ResNet-18 on a 5-class flower dataset for 3 epochs; show the confusion matrix improving after each epoch.

Discussion: Give students an imbalanced dataset and ask why 92% accuracy is a trap—lead them to recall and F1.

Pitfall to demo: Deliberately apply softmax before CrossEntropyLoss and let students watch training stall, then fix it.

Recap Image classification is the foundation of computer vision: one image in, one label out. The backbone you fine-tune here is the same backbone that powers detection and segmentation—and the Module 7.3 architectures like ResNet and ViT.