Flatten
- Machine Learning course
- Advanced
- Lesson 45/78
- 5 min
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 callIn networks using global average pooling this problem never comes up: the linear layer's input is always the channel count.
Practice 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.
Starter code
import torch
from torch import nn
# buraya yazSample 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()))Expected output
torch.Size([4, 10]) 9098
Questions and answers
Review the lesson questions together with their answers. To test yourself, solve the lesson interactively.
1 What shape does a (32, 64, 5, 5) output take after <code>nn.Flatten()</code>?
- (32, 320)
- (32, 1600) Correct answer
- (51200,)
- (32, 64)
Explanation: 64 · 5 · 5 = 1600; the batch dimension is kept.
2 Flatten the axes other than the batch.
nn.____()nn.Flatten()Explanation: nn.Flatten uses start_dim=1 by default.
3 Match the sizes.
| Term | Meaning |
|---|---|
(N, 16, 7, 7) |
784 features |
(N, 8, 13, 13) |
1352 features |
nn.LazyLinear |
Finds the size itself |
AdaptiveAvgPool2d(1) |
Features = channels |