Building an Image Classifier

Build the mental model

A working image classifier almost always follows the same two-stage shape: a feature extractor made of alternating Conv2d, activation, and pooling layers that progressively compresses a raw image into a compact set of learned features, followed by a classifier head made of one or more Linear layers that maps those features to a score (logit) for each possible class. The feature extractor's job is representation — turning raw pixels into abstract, discriminative features — while the classifier head's job is decision — combining those features into a final judgment. Between the two sits a `Flatten` step, since conv layers output a multi-dimensional (channels, height, width) tensor per image but Linear layers expect a flat vector per sample. The output layer has no activation function of its own inside the model; it produces raw logits, one per class, and the loss function (typically cross-entropy) internally applies softmax to turn those logits into class probabilities during training. This same feature-extractor-plus-head pattern scales from tiny toy CNNs up to today's largest vision models, which just add far more, and more sophisticated, layers to each stage.

Connect it to a real scenario

Building the duplicate/low-quality lesson image detector means assembling exactly this shape: a couple of Conv2d+ReLU+MaxPool blocks compress each uploaded screenshot down to a small feature map capturing texture, blur, and layout cues, then a Flatten and one or two Linear layers turn that into a score for each output category — 'sharp original', 'blurry duplicate', 'watermarked'. Because the whole model is just one `nn.Module`, the platform can save it, load it into its image-upload pipeline, and run a forward pass on every new screenshot the moment a course author uploads it, flagging low-quality or duplicate images before they ever reach a published lesson.

Almost every image CNN follows the same shape: conv+pool blocks extract features, then a Flatten and Linear layers classify them.
The final output layer produces raw logits (one per class) with no built-in activation — the loss function applies softmax internally.
A `Flatten` layer is needed between the conv feature extractor and the Linear classifier head to turn a multi-dimensional tensor into a flat vector.

Try the working example

python

import torch
import torch.nn as nn

class SmallCNN(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        self.block1 = nn.Sequential(
            nn.Conv2d(3, 8, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
        )
        self.block2 = nn.Sequential(
            nn.Conv2d(8, 16, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
        )
        self.flatten = nn.Flatten()
        self.fc1 = nn.Linear(16 * 8 * 8, 32)
        self.fc2 = nn.Linear(32, num_classes)

    def forward(self, x):
        x = self.block1(x)   # 32x32 -> 16x16
        x = self.block2(x)   # 16x16 -> 8x8
        x = self.flatten(x)
        x = torch.relu(self.fc1(x))
        return self.fc2(x)

model = SmallCNN(num_classes=10)

# Fake batch of 4 RGB "images", 32x32 pixels
images = torch.randn(4, 3, 32, 32)
logits = model(images)
print("Output shape:", logits.shape)

You should see

It prints `Output shape: torch.Size([4, 10])` — one row of 10 class logits for each of the 4 images in the batch.

5-minute try-it

Add a `nn.Softmax(dim=1)` call to the logits after the forward pass and print the result to confirm each row now sums to 1.0.

One important caution

Getting the Linear layer's input size wrong (e.g. writing `nn.Linear(16*16*16, 32)` when the feature map is actually 8x8, not 16x16) — this raises a matrix-shape mismatch error at the first forward pass since it doesn't match the flattened tensor size.

Applying `softmax` inside the model's forward pass before using `nn.CrossEntropyLoss` for training — that loss function already applies softmax internally, so doing it twice distorts the gradients and hurts training.

PyTorch Docs — Training a Classifier (CIFAR-10) — Deep Learning