Building CNNs in PyTorch

By the end of this lesson, you will be able to...

• translate a small CNN plan into PyTorch modules;

• predict the tensor shape after each stage and diagnose a broken classifier head;

• compare flattening with adaptive average pooling under a parameter and location-detail constraint.

Idea in one sentence: A CNN implementation is trustworthy when its modules, tensor shapes, output meaning, and train/eval behavior all match the model promise.

Core Insight

Suppose the plant-inspection team now needs its leaf classifier to run. Their architecture sketch says: two convolution blocks, two reductions in resolution, then two disease logits. The first implementation crashes at Linear: it receives 8,192 values per image, but the layer was built for 2,048.

The naive model says that the convolutional blocks are correct because they ran. The error is evidence that a CNN is not only a diagram. Each block makes a concrete tensor promise to the next one. The stronger model is a shape ledger: write the batch, channel, height, and width after every operation before trusting the classifier head.

This lab builds one bounded classifier. The code is not a deployment system and does not prove model quality. It makes the architecture from the previous lessons executable and inspectable.

The Contract of the Small Model

Assume an input batch of RGB leaf images resized and normalized by an already-declared preprocessing path:

input:  (B, 3, 64, 64)
output: (B, 2) raw logits for [healthy, diseased]

The model returns logits, not probabilities. A training loss such as cross entropy expects those raw scores; a later inference boundary can convert them to probabilities and apply a decision rule. Keeping that distinction visible prevents a common error: adding a softmax merely because the output is called a classifier.

The Plan Before the Code

The feature extractor will preserve 64 × 64 long enough to learn small marks, then pool twice. Padding 1 keeps a 3 × 3 stride-1 convolution at the same spatial size. The numbers are a teaching design, not a claim that 16 and 32 channels are optimal.

The initial head is tempting:

Flatten (B, 32, 16, 16) -> (B, 8192) -> Linear(8192, 2)

It works only while the input resolution and every prior stride remain exactly as expected. It also gives the head 16,386 parameters including biases. That may be acceptable in a stated design, but it is a choice, not an invisible default.

The Better Boundary: Pool Before the Head

For this leaf-level classification promise, use AdaptiveAvgPool2d((1, 1)) before flattening. It summarizes each of the 32 feature channels over its spatial positions:

(B, 32, 16, 16)
-> adaptive average pool
-> (B, 32, 1, 1)
-> flatten
-> (B, 32)
-> Linear(32, 2)

The final linear layer now has 66 parameters including biases. This buys a much smaller head and makes the head independent of the incoming spatial size. It costs exact location information: it is a better fit for “is disease present?” than “where is the lesion boundary?” The previous lesson supplies the reason; the code merely makes the boundary real.

A Complete, Small Module

import torch
from torch import nn


class LeafCNN(nn.Module):
    def __init__(self, classes=2):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 16, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(16, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
        )
        self.head = nn.Sequential(
            nn.AdaptiveAvgPool2d((1, 1)),
            nn.Flatten(),
            nn.Linear(32, classes),
        )

    def forward(self, x):
        return self.head(self.features(x))


model = LeafCNN()
x = torch.randn(8, 3, 64, 64)  # synthetic batch for a shape check
logits = model(x)
assert logits.shape == (8, 2)

The assert tests only the declared shape. It does not show that labels are correct, preprocessing matches validation, or the model has learned useful features. Those are different claims and need their own evidence.

So far, we have seen that the module follows the ledger exactly. This matters because a shape check catches an implementation boundary error before an opaque training failure hides it.

Trace a Failure Instead of Guessing

Consider this accidental edit: the second convolution becomes nn.Conv2d(16, 48, ...), but the head stays nn.Linear(32, 2) after adaptive pooling.

features output: (B, 48, 16, 16)
pooled output:   (B, 48, 1, 1)
flattened:       (B, 48)
head expects:    (B, 32)

The failure is not “PyTorch is confused.” The producer and consumer disagree about channels. The repair is to change the head to nn.Linear(48, 2), or to restore the block's 32 output channels. Print or assert shapes at the boundary where a module hands data to another module; do not change dimensions until the error disappears by accident.

Another common wrong turn is treating an input shaped (B, 64, 64, 3) as if it were PyTorch image input. Conv2d normally expects (B, C, H, W), so channels belong second. The conversion belongs at one declared input boundary, not scattered through model code.

Train and Evaluate Are Different States

This tiny module has no BatchNorm or Dropout, so train() and eval() produce the same layer behavior here. That is an assumption of the example, not a general CNN rule. Once a model contains Dropout or BatchNorm, training and inference modes differ. Normal inference should explicitly use model.eval() and torch.no_grad().

model.eval()
with torch.no_grad():
    logits = model(x)

This does not replace a validation protocol or an inference contract. It prevents the framework from retaining gradients during ordinary inference and sets mode-sensitive modules to their evaluation behavior.

Costs, Limits, and Signals

The trade-off is inspectable. Global-style average pooling reduces head parameters and eases resolution changes, but it can discard evidence about a lesion's exact arrangement. Flattening retains more spatial layout for the head, but can create a large resolution-dependent parameter block. Choose according to the output promise and compare validation slices, latency, memory, and parameter count.

Useful signals are: the shape ledger for correctness; a one-batch forward assertion for module boundaries; the number of parameters in the head; and accuracy broken out by small lesions or shifted leaves. A successful forward pass is necessary but not evidence that the model generalizes.

This lab also does not solve data loading, augmentation policy, checkpoint selection, serving, or fleet monitoring. The next cluster changes the architectural problem from spatial structure to sequence state.

Common Confusions

Confusion: “Adaptive pooling makes any input resolution safe.”

Why it is tempting: the final linear layer keeps the same input width after the pooling operation. Better model: it stabilizes one boundary in the head. Earlier convolutions still perform more work on a larger image, and changing resolution can change which visual detail is visible to the model. Test the revised input contract rather than treating the shape assertion as a quality result.

Confusion: “A forward pass proves the classifier is ready.”

Why it is tempting: the most visible implementation error has disappeared. Better model: a forward pass proves only that the declared tensor boundary was satisfied for that batch. The labels, preprocessing, training evidence, output decoding, and evaluation split still determine whether the scores support a decision.

Confusion: “Changing model.eval() repairs a shape mismatch.”

Why it is tempting: both are common lines in an inference snippet. Better model: evaluation mode controls modules whose behavior depends on mode; it does not reorder channels, change a convolution's output width, or repair a wrong Linear input size. Diagnose the contract that actually failed.

The fastest useful debugging loop is therefore: state the expected shape, run one synthetic batch, inspect the first mismatch, and repair only that boundary.

Check Your Understanding

Check: What is the shape after MaxPool2d(2) receives (B, 16, 64, 64) with its usual stride?

Think first, then reveal.

Answer: (B, 16, 32, 32). Pooling keeps the channel count and halves both spatial dimensions.

Check: A feature block now emits (B, 48, 16, 16). What must a head using adaptive average pooling expect before its final two-class linear layer?

Answer: It receives (B, 48) after pooling and flattening, so the layer must be Linear(48, 2).

Practice: Repair the Boundary

You change the input contract from 64 × 64 to 128 × 128. The model uses two pools and an adaptive-average-pooling head. Which shapes change, and which head parameter does not?

Model answer: The feature maps become (B, 16, 128, 128), (B, 16, 64, 64), (B, 32, 64, 64), and (B, 32, 32, 32). Adaptive pooling still produces (B, 32, 1, 1), so Linear(32, 2) does not change. This preserves the head boundary, but the increased spatial maps cost more compute and memory; validate that the revised resolution helps the real task.

Resources

• [DOCUMENTATION] PyTorch: Conv2d — Focus: channel and spatial-shape conventions.
• [DOCUMENTATION] PyTorch: MaxPool2d — Focus: pooling output shape.
• [DOCUMENTATION] PyTorch: AdaptiveAvgPool2d — Focus: a resolution-tolerant classifier-head boundary.
• [TUTORIAL] PyTorch: Training a Classifier — Focus: connecting a bounded model to a training and evaluation loop.

Key Takeaways

• A CNN module is a chain of tensor contracts; a shape ledger makes those contracts inspectable.
• Adaptive average pooling can simplify a classifier head, but it trades away spatial arrangement and must match the task promise.
• Diagnose a CNN by tracing the producer and consumer shapes at each boundary, then test generalization separately from “the code runs.”