Assemble convolution, pooling, and dense layers into a working image classifier. See the standard pattern — conv blocks that extract features, then a head that decides.
Why: almost every CNN follows the same recipe — stacked conv+ReLU+pool blocks that extract increasingly abstract features, then a flatten and dense head that maps them to class scores. When: use this template as your starting point for any image task. Where: spatial size shrinks while channel depth grows as you go deeper.
INPUT image
|
[ Conv -> ReLU -> Pool ] x N feature extractor
| (H,W shrink; channels grow)
Flatten
|
[ Linear -> ReLU -> Linear ] classifier head
|
class scores (logits)Why: writing the recipe as an nn.Module gives you a real, trainable image classifier — feature extractor plus head. When: adapt the channel counts and final output size to your dataset. Where: the flattened size feeding the first linear layer depends on the input size and how much pooling reduced it — compute it or use nn.LazyLinear.
import torch, torch.nn as nn
class SmallCNN(nn.Module):
def __init__(self, num_classes=10):
super().__init__()
self.features = nn.Sequential(
nn.Conv2d(3, 16, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2), # 32->16
nn.Conv2d(16, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2), # 16->8
)
self.head = nn.Sequential(
nn.Flatten(),
nn.Linear(32 * 8 * 8, 128), nn.ReLU(),
nn.Linear(128, num_classes),
)
def forward(self, x):
return self.head(self.features(x))
model = SmallCNN()
print(model(torch.randn(4, 3, 32, 32)).shape) # [4, 10] logitsWhy: a CNN trains with the exact same loop as the MLP from the PyTorch course — the architecture changed, the training procedure did not. When: use CrossEntropyLoss, an Adam optimizer, and DataLoaders of image batches. Where: image data needs normalization and augmentation (flips, crops) applied to the training set only.
import torch, torch.nn as nn
model = SmallCNN(num_classes=10)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
for epoch in range(10):
model.train()
for images, labels in train_loader: # batches of (N,3,32,32)
optimizer.zero_grad()
loss = criterion(model(images), labels)
loss.backward()
optimizer.step()
# Same five-step loop as the PyTorch course — only the model differs.