Flattening reshapes a multi-dimensional feature map tensor into a single flat vector โ the standard bridge between a CNN's spatial convolutional layers and a final fully-connected (dense) classification head.
What Flattening Does
A 3-D feature map tensor is reshaped into a 1-D vector, simply by laying out every value in a fixed order โ no values are lost or changed, only the shape. Fully-connected layers (as covered throughout the Neural Network Fundamentals category) expect a plain vector input; flattening is what makes a CNN's spatial output compatible with that expectation.
Numerical Example
A feature map of shape \((2, 2, 2)\) โ 2 channels, each a 2ร2 grid: channel 1 = \(\begin{bmatrix}1&2\\3&4\end{bmatrix}\), channel 2 = \(\begin{bmatrix}5&6\\7&8\end{bmatrix}\). Flattened (in PyTorch's default row-major order): \([1,2,3,4,5,6,7,8]\) โ a single vector of length 8, exactly \(2\times2\times2\).
Code
import torch
import torch.nn as nn
x = torch.tensor([[[[1.,2.],[3.,4.]],
[[5.,6.],[7.,8.]]]]) # shape (1, 2, 2, 2) -- batch, channels, H, W
flatten = nn.Flatten()
output = flatten(x)
print(output.shape) # torch.Size([1, 8])
print(output) # tensor([[1., 2., 3., 4., 5., 6., 7., 8.]])
Where Flattening Fits in a Full CNN
model = 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),
nn.Flatten(), # bridges spatial feature maps to a flat vector
nn.Linear(32 * 8 * 8, 10) # the flat vector's length MUST match this Linear layer's input size exactly
)
Common Mistakes
- Forgetting to include a flattening step before a fully-connected layer entirely โ this raises a clear shape-mismatch error, since
nn.Linearexpects a 2-D input (batch, features), not a 4-D feature map tensor. - Hard-coding a following
nn.Linearlayer's expected input size incorrectly after changing an earlier layer's channel count, kernel size, stride, or padding โ any of these changes the flattened vector's length, and the fully-connected layer's declared input size must be recomputed to match exactly.
Interview Relevance
Q: "Why does a CNN need a flattening step before its final classification layers?" Convolutional and pooling layers output multi-dimensional spatial tensors (channels ร height ร width), but fully-connected layers expect a plain 1-D vector input per example. Flattening reshapes the spatial output into exactly that vector form, preserving every value while making the shape compatible with the dense layers that follow.
Practice Question
A CNN's final convolutional block outputs feature maps of shape \((64, 4, 4)\) per example. What length vector does flattening produce, and what must the following nn.Linear layer's input size be set to?