Background
Sections
IntroductionModule 01 β€” Tensors 🧊01 Β· Creating Tensors 🧊02 Β· Indexing & Reshaping πŸ”ͺ03 Β· Tensor Math βž—04 Β· Device Placement πŸ–₯️⚑Module 02 β€” Autograd βš™οΈ01 Β· The Computational Graph πŸ•ΈοΈ02 Β· Backward Pass & Gradients ⬅️03 Β· Turning Autograd Off πŸ›‘04 Β· Gradient Gotchas πŸͺ€Module 03 β€” Neural Networks01 Β· The `nn.Module` Basics 🧱02 Β· Common Layers 🧩03 Β· Building a Network πŸ—οΈ04 Β· Inspecting Models πŸ”Module 04 β€” Data Handling πŸ—‚οΈ01 Β· Dataset Basics πŸ“‡02 Β· The DataLoader 🚚03 Β· Transforms 🎨04 Β· Splits & Built-in Datasets βœ‚οΈModule 05 β€” The Training Loop πŸ”01 Β· Loss Functions 🎯02 Β· Optimizers βš™οΈ03 Β· The Training Loop πŸ”04 Β· Evaluation & Metrics πŸ“ŠModule 06 β€” Saving & Loading πŸ’Ύ01 Β· `state_dict` Basics πŸ’Ύ02 Β· Checkpoints & Resuming ⏸️03 Β· Loading for Inference πŸš€04 Β· Best Model & Early Stopping πŸ…Module 07 β€” Computer Vision πŸ‘οΈ01 Β· Convolutions πŸ”²02 Β· CNN Architecture πŸ›οΈ03 Β· Transfer Learning πŸ”04 Β· Image Classification Project πŸ§ͺModule 08 β€” NLP & Transformers πŸ’¬01 Β· Text Data & Tokenization πŸ”€02 Β· Embeddings 🧭03 Β· Recurrent Layers & LSTMs πŸ”„04 Β· Intro to Transformers ⚑Module 09 β€” Ecosystem: PyTorch Lightning ⚑01 Β· Why Lightning? πŸ€”02 Β· The LightningModule 🧩03 Β· The Trainer πŸŽ›οΈ04 Β· DataModules & Callbacks 🧰Module 10 β€” Deployment & Optimization πŸš€01 Β· Exporting Models πŸ“¦02 Β· `torch.compile` ⚑03 Β· Inference Optimization πŸͺΆ04 Β· Serving & Next Steps πŸŽ“

02 Β· CNN Architecture πŸ›οΈ

5 min read

One convolution detects patterns. A CNN stacks convolutions, pooling, and a classifier into a machine that turns raw pixels into a decision.


1. The 'Why'

A single convolution (sub-module 01) finds local patterns, but recognizing a real object takes more: you need to combine simple features into complex ones, shrink the spatial resolution so the network can "see" larger structures without exploding in cost, and finally collapse everything into a class prediction. A CNN architecture is the recipe for arranging these steps. The classic pattern β€” repeat (convolution β†’ activation β†’ pooling) a few times to build up rich, spatially-compressed feature maps, then flatten and pass through Linear layers to produce class scores β€” has powered image recognition for over a decade and remains the mental model behind even the fanciest modern networks.

The new ingredient here is pooling, and understanding why it exists is the key to the module. As convolutions detect features, nn.MaxPool2d periodically downsamples the feature maps β€” keeping the strongest response in each little neighborhood and discarding the rest. This does three valuable things at once: it shrinks the data (so deeper layers are cheaper), it grows each neuron's "receptive field" (so later layers respond to larger regions of the original image), and it adds a bit of robustness to small shifts. Combine convolutions for detecting, pooling for compressing, and linear layers for deciding, and you have a complete CNN β€” the architecture you'll train in the capstone and adapt via transfer learning next.


2. Core Concepts

nn.MaxPool2d(kernel_size). Slides a window over each feature map and keeps only the maximum value in each window, downsampling the spatial dimensions. A 2Γ—2 max-pool halves height and width. It has no learnable parameters.

The canonical CNN shape. A feature extractor (stacked conv β†’ ReLU β†’ pool blocks) followed by a classifier head (flatten β†’ Linear β†’ ReLU β†’ Linear). The extractor learns what's in the image; the head turns features into class scores.

Flattening the bridge. Between the last conv/pool and the first Linear, you must flatten the (batch, channels, H, W) feature maps into (batch, features). Use nn.Flatten() or x.view(x.size(0), -1).

Computing the flattened size. The first Linear's in_features equals channels Γ— H Γ— W after the final pool. Get this wrong and you get a shape error β€” trace the sizes through the network (or use nn.LazyLinear / nn.AdaptiveAvgPool2d to sidestep it).


3. Code in Action

Max pooling in isolation

import torch
import torch.nn as nn

pool = nn.MaxPool2d(kernel_size=2)     # 2x2 window -> halves H and W
x = torch.randn(1, 16, 32, 32)         # 16 feature maps, 32x32
print(pool(x).shape)                   # torch.Size([1, 16, 16, 16]) -> spatially halved

A complete CNN (custom nn.Module)

import torch
import torch.nn as nn

class SimpleCNN(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=3, padding=1),  # 1x28x28 -> 32x28x28
            nn.ReLU(),
            nn.MaxPool2d(2),                             # -> 32x14x14
            nn.Conv2d(32, 64, kernel_size=3, padding=1), # -> 64x14x14
            nn.ReLU(),
            nn.MaxPool2d(2),                             # -> 64x7x7
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),                    # 64x7x7 -> 3136 features
            nn.Linear(64 * 7 * 7, 128),      # in_features MUST match 64*7*7
            nn.ReLU(),
            nn.Dropout(0.25),                # regularization (Module 03)
            nn.Linear(128, num_classes),     # -> raw logits per class
        )

    def forward(self, x):
        x = self.features(x)                 # extract spatial features
        return self.classifier(x)            # decide the class

model = SimpleCNN()
out = model(torch.randn(8, 1, 28, 28))       # a batch of 8 grayscale 28x28 images
print(out.shape)                             # torch.Size([8, 10]) -> logits

Avoiding the flatten-size headache with adaptive pooling

import torch.nn as nn

# AdaptiveAvgPool2d forces a fixed output size regardless of input dimensions,
# so the following Linear always sees the same feature count.
head = nn.Sequential(
    nn.AdaptiveAvgPool2d((1, 1)),   # collapse each feature map to a single number
    nn.Flatten(),                   # (batch, channels, 1, 1) -> (batch, channels)
    nn.Linear(64, 10),              # in_features = channels only β€” no H,W arithmetic
)

4. Common Pitfalls

Pitfall 1 β€” Mis-computing the first Linear's in_features. This is the #1 CNN bug. After your conv/pool stack, the feature volume is channels Γ— H Γ— W; that exact product must be the in_features of the first linear layer. Trace the spatial size through each pool, or use AdaptiveAvgPool2d/LazyLinear to avoid the arithmetic.

Pitfall 2 β€” Forgetting to flatten before Linear. nn.Linear expects (batch, features), but conv output is 4-D (batch, C, H, W). Insert nn.Flatten() between the extractor and the classifier, or you'll get a dimension error.

Pitfall 3 β€” Pooling away too much, too fast. Aggressive pooling early (or too many pools) can shrink the feature maps to 1Γ—1 before the network has learned useful features, starving the classifier. Balance the number of pools against the input resolution.


5. Further Reading & Watch List


⬅️ Prev: 01 Β· Convolutions Β· ➑️ Next: 03 Β· Transfer Learning