PyTorch Pooling Layers

Pooling layers reduce the spatial dimensions (height and width) of feature maps. They shrink the data, reduce computation in later layers, and make the network less sensitive to the exact position of features. A cat's ear detected two pixels to the left should still be recognized as a cat's ear — pooling helps achieve this positional flexibility.

Why Pooling Exists

Without pooling:
  Conv layers produce large feature maps
  Downstream layers process enormous amounts of data
  Network becomes very slow and prone to overfitting

With pooling after each conv block:
  Spatial size halves at each step
  Computation drops by 4× per pool layer
  Network forces itself to keep the most important information

MaxPool2d — Keep the Maximum

MaxPool2d slides a window over the feature map and keeps only the maximum value in each window. If a strong feature (like an edge) exists anywhere in the pooling window, it gets preserved. Weak activations get discarded.

import torch
import torch.nn as nn

pool = nn.MaxPool2d(kernel_size=2, stride=2)   # 2×2 window, move 2 steps

x = torch.tensor([[[[1.0, 3.0, 2.0, 4.0],
                    [5.0, 6.0, 7.0, 8.0],
                    [2.0, 1.0, 4.0, 3.0],
                    [9.0, 2.0, 1.0, 6.0]]]])

out = pool(x)
print(out)
# tensor([[[[6., 8.],
#           [9., 6.]]]])

MaxPool Diagram

Input 4×4:          Pool 2×2, stride 2 → Output 2×2:
┌──┬──┬──┬──┐
│1 │3 │2 │4 │       max(1,3,5,6)=6   max(2,4,7,8)=8
├──┼──┼──┼──┤    →  ┌──┬──┐
│5 │6 │7 │8 │       │6 │8 │
├──┼──┼──┼──┤       ├──┼──┤
│2 │1 │4 │3 │       │9 │6 │
├──┼──┼──┼──┤       └──┴──┘
│9 │2 │1 │6 │       max(2,1,9,2)=9   max(4,3,1,6)=6
└──┴──┴──┴──┘

AvgPool2d — Keep the Average

AvgPool2d takes the average value in each window instead of the maximum. It produces smoother output but may dilute strong signals. Some architectures use average pooling in their final global pooling layer.

avg_pool = nn.AvgPool2d(kernel_size=2, stride=2)
out = avg_pool(x)
print(out)
# average of each 2×2 block

Global Average Pooling

Global Average Pooling collapses an entire feature map (any spatial size) to a single value per channel by averaging all spatial positions. It completely removes the spatial dimensions. This replaces the flatten + fully connected approach used in older CNNs and produces models that accept any input size.

global_avg = nn.AdaptiveAvgPool2d((1, 1))   # output always 1×1

feature_map = torch.rand(8, 64, 14, 14)    # batch, channels, H, W
out = global_avg(feature_map)
print(out.shape)                            # torch.Size([8, 64, 1, 1])

out = out.squeeze(-1).squeeze(-1)           # remove 1×1 spatial dims
print(out.shape)                            # torch.Size([8, 64])

Global Average Pooling diagram:

Feature map [8, 64, 14, 14]:
  64 channels, each 14×14 = 196 values
         ↓ average each 14×14 map to a single number
Output [8, 64]:
  64 channels, each now a single value
  → directly feed to a Linear(64, num_classes) layer

AdaptiveAvgPool2d — Output Any Size

The adaptive version lets you specify the output size and PyTorch figures out the window size automatically. This makes it easy to fix the output size regardless of the input image dimensions.

pool = nn.AdaptiveAvgPool2d((4, 4))   # output always 4×4

x1 = torch.rand(1, 32, 28, 28)
x2 = torch.rand(1, 32, 56, 56)

print(pool(x1).shape)   # torch.Size([1, 32, 4, 4])
print(pool(x2).shape)   # torch.Size([1, 32, 4, 4])  ← same output size!

MaxPool vs AvgPool — When to Use Each

MaxPool2d:
  → Preserves the strongest detected features
  → Most common in classification CNNs (VGG, early AlexNet)
  → Good when you care about whether a feature exists

AvgPool2d:
  → Smooths the feature map
  → Used in global pooling (ResNet, EfficientNet)
  → Good when you care about feature density across a region

Spatial Size After Pooling

Input:  [batch, channels, H, W]
MaxPool2d(2, stride=2):
Output: [batch, channels, H/2, W/2]

Example sizes through a CNN:
  Input:         [8, 3, 224, 224]
  After Conv+BN: [8, 32, 224, 224]
  After Pool:    [8, 32, 112, 112]   ← halved
  After Conv+BN: [8, 64, 112, 112]
  After Pool:    [8, 64, 56, 56]     ← halved again
  After Conv+BN: [8, 128, 56, 56]
  After Pool:    [8, 128, 28, 28]
  → GlobalAvgPool: [8, 128]
  → Linear(128, 10): [8, 10]

1D and 3D Pooling

PyTorch also provides MaxPool1d for sequence data (NLP, audio) and MaxPool3d for video or 3D medical imaging. The same concepts apply — the kernel slides across the relevant number of dimensions.

Summary

Pooling layers reduce spatial dimensions while retaining important features. MaxPool2d keeps the maximum value in each window — perfect for detecting whether a feature is present. AvgPool2d averages the values — better for global representations. AdaptiveAvgPool2d((1,1)) implements global average pooling, collapsing each feature map to a single number and removing spatial dimensions entirely. Use global pooling to make your network accept variable-size inputs. Pooling reduces computation, controls overfitting, and builds positional flexibility into the network.