Machine Learning
Flatten
Flatten reduces multi-dimensional feature maps to one dimension.
Flatten
The output of the convolutional body has shape (N, C, H, W); the classifier layer wants (N, features). nn.Flatten() merges all axes except the batch dimension:
model = nn.Sequential(
nn.Conv2d(1, 8, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2), # 28 -> 14
nn.Conv2d(8, 16, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2), # 14 -> 7
nn.Flatten(), # 16·7·7 = 784
nn.Linear(16 * 7 * 7, 10))
Giving the body a dummy input is a practical way to work out the linear layer's input size.
Finding the size automatically
When the input size changes, computing the linear layer's size by hand is error-prone. Two common fixes: measuring with a dummy input, or using nn.LazyLinear:
govde = nn.Sequential(nn.Conv2d(1, 8, 3), nn.ReLU(), nn.MaxPool2d(2), nn.Flatten())
with torch.no_grad():
n = govde(torch.zeros(1, 1, 28, 28)).shape[1]
print(n) # 8·13·13 = 1352
bas = nn.LazyLinear(10) # learns its input size on the first call
In networks using global average pooling this problem never comes up: the linear layer's input is always the channel count.
Exercise A small CNN
Build the two-block model above (1→8→16 channels, padding 1, MaxPool2d(2) in each block, then Flatten and nn.Linear(16 * 7 * 7, 10)). Print the output shape for a (4, 1, 28, 28) input and the total parameter count on one line.
Show example solution
import torch
from torch import nn
model = nn.Sequential(
nn.Conv2d(1, 8, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
nn.Conv2d(8, 16, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
nn.Flatten(),
nn.Linear(16 * 7 * 7, 10))
print(model(torch.zeros(4, 1, 28, 28)).shape, sum(p.numel() for p in model.parameters()))
Test yourself
-
What shape does a (32, 64, 5, 5) output take after <code>nn.Flatten()</code>?
- (32, 320)
- (32, 1600)
- (51200,)
- (32, 64)
Show answer
(32, 1600) — 64 · 5 · 5 = 1600; the batch dimension is kept.
-
Flatten the axes other than the batch.
nn.____()
Show answer
nn.Flatten()
-
Match the sizes.
- (N, 16, 7, 7)
- 784 features
- (N, 8, 13, 13)
- 1352 features
- nn.LazyLinear
- Finds the size itself
- AdaptiveAvgPool2d(1)
- Features = channels
Next lesson: Summary - CNN
