Building an Image Classifier

Build the mental model

An image classifier is simply a feature extractor bolted onto a decision head. Each Conv2d layer detects local patterns (edges, textures, shape fragments); ReLU adds non-linearity; MaxPool shrinks the spatial dimensions while keeping the strongest activations. Stacking blocks lets the network build progressively more abstract features — from low-level edges to high-level object parts. At the end, Flatten turns the spatial feature map into a vector, and a Linear layer (the classifier head) outputs one score per class. This beats a plain fully-connected network on images because convolution shares weights across the whole spatial extent — the same small filter slides everywhere, so you don't need a separate weight for every pixel position.

The training loop itself is the same forward → loss → backward → optimizer.step() pattern from the PyTorch fundamentals course, but here the input is an image tensor and cross-entropy loss implicitly converts the raw class-score logits into probabilities before comparing them against the true integer labels. Training repeatedly on one fixed batch (deliberately overfitting it) is a useful sanity check for a new architecture or training loop: if the loss refuses to go down even on a single batch, something in the model or loop is broken, and you want to know that before scaling up to a real dataset.

Connect it to a real scenario

Suppose the Tutorial Platform is prototyping a feature that auto-tags instructor-uploaded course thumbnails by topic category (Programming, Design, Math, and so on). Before touching any real thumbnails, you'd build exactly this kind of small CNN and verify the training loop works correctly on synthetic data first — confirming the architecture and loop are bug-free before ever pointing it at the real thumbnail dataset.

• A CNN extracts features with Conv2d+ReLU+MaxPool blocks, then Flatten+Linear turns them into class scores.
• The training loop pattern is identical to the PyTorch fundamentals course — only now the input is an image batch.
• Overfitting a single fixed batch is a fast way to sanity-check that a new architecture and training loop are wired correctly.

Try the working example

import torch
import torch.nn as nn
import torch.optim as optim

torch.manual_seed(0)

class SimpleCNN(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 16, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),  # 32x32 -> 16x16
            nn.Conv2d(16, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),  # 16x16 -> 8x8
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(32 * 8 * 8, num_classes),
        )

    def forward(self, x):
        x = self.features(x)
        return self.classifier(x)

model = SimpleCNN(num_classes=10)
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
criterion = nn.CrossEntropyLoss()

# Fake batch: 8 RGB images, 32x32, with random integer labels 0-9
images = torch.randn(8, 3, 32, 32)
labels = torch.randint(0, 10, (8,))

for step in range(5):
    optimizer.zero_grad()
    outputs = model(images)
    loss = criterion(outputs, labels)
    loss.backward()
    optimizer.step()
    print(f"step {step}: loss = {loss.item():.4f}")

Prints five lines in the form 'step 0: loss = X.XXXX' through 'step 4: loss = X.XXXX'. Because the model trains repeatedly on the exact same fixed batch, the loss should trend downward across the five steps overall (it may not be strictly monotonic step-to-step), starting near ln(10) ≈ 2.30 — the expected cross-entropy loss for a freshly initialized 10-way classifier making essentially random guesses.

PyTorch Tutorials — Training a Classifier (CIFAR10) — Computer Vision