PyTorch Pooling

Pooling shrinks a feature map by replacing each small patch with one value. Max pooling keeps the largest value in each patch.

One patch

A 2 by 2 patch pools down to its largest value. The other three values are gone.

patch = torch.tensor([[[[0.0, 2.0], [1.0, 0.0]]]])
print(nn.MaxPool2d(2)(patch)[0, 0])

Position is lost

Two patches with the winner in different corners pool to the same value. The output cannot say where the winner sat.

a = torch.tensor([[[[2.0, 0.0], [0.0, 0.0]]]])
b = torch.tensor([[[[0.0, 0.0], [0.0, 2.0]]]])
pool = nn.MaxPool2d(2)
print(pool(a).item(), pool(b).item(), pool(a).item() == pool(b).item())

Average pooling

Average pooling keeps the mean of the patch instead. It softens the signal rather than picking a winner.

patch = torch.tensor([[[[0.0, 2.0], [1.0, 1.0]]]])
print(nn.AvgPool2d(2)(patch)[0, 0])

Example

import torch
from torch import nn

feature_map = torch.tensor([[
    [0.0, 2.0, 1.0, 0.0, 3.0, 1.0],
    [1.0, 0.0, 0.0, 2.0, 1.0, 0.0],
    [2.0, 5.0, 1.0, 1.0, 0.0, 4.0],
    [0.0, 1.0, 3.0, 0.0, 2.0, 0.0],
    [1.0, 0.0, 4.0, 2.0, 0.0, 1.0],
    [0.0, 3.0, 1.0, 0.0, 1.0, 2.0],
]]).unsqueeze(0)

pool = nn.MaxPool2d(2)
pooled = pool(feature_map)

print("Before:")
print(feature_map[0, 0])
print("After:")
print(pooled[0, 0])
print("Cells before:", feature_map.numel())
print("Cells after:", pooled.numel())

Key points

Feature map: a grid of scores left by one kernel
Use MaxPool2d(2) to halve the height and width of a map
Pooling keeps the strongest score in each patch
Pooling throws away the other values in the patch
Pooling also throws away the exact position of the winning value

In this example, a 6 by 6 feature map holds scores with a few strong cells. MaxPool2d(2) splits the map into 2 by 2 patches. Each patch prints as its single largest value in the pooled map. The strong 5 and 4 survive, and weaker neighbours vanish. The cell count drops from 36 to 9, which is the honest price of the smaller map. The pooled map still shows where strong signals fired, at a coarser scale.

Note

Pooling has no weights, so it has no parameters to learn

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

import torch
from torch import nn

feature_map = torch.tensor([[
    [0.0, 2.0, 1.0, 0.0, 3.0, 1.0],
    [1.0, 0.0, 0.0, 2.0, 1.0, 0.0],
    [2.0, 5.0, 1.0, 1.0, 0.0, 4.0],
    [0.0, 1.0, 3.0, 0.0, 2.0, 0.0],
    [1.0, 0.0, 4.0, 2.0, 0.0, 1.0],
    [0.0, 3.0, 1.0, 0.0, 1.0, 2.0],
]]).unsqueeze(0)

pool = nn.MaxPool2d(2)
pooled = pool(feature_map)

print("Before:")
print(feature_map[0, 0])
print("After:")
print(pooled[0, 0])
print("Cells before:", feature_map.numel())
print("Cells after:", pooled.numel())

Exercise

What does MaxPool2d return for the patch [[1, 5], [3, 2]]?

5 3 2.75 11