PyTorch Global Average Pooling

Global average pooling averages each whole feature map down to one number. The result is one score per map, not one score per cell.

One map, one number

A 2 by 2 map averages to its mean. The four values become a single 2.

m = torch.tensor([[[[1.0, 3.0], [2.0, 2.0]]]])
print(nn.AdaptiveAvgPool2d(1)(m)[0, 0])

Every map shrinks

Eight 4 by 4 maps pool to eight single cells. The channel count stays at eight.

maps = torch.rand(1, 8, 4, 4)
print(nn.AdaptiveAvgPool2d(1)(maps).shape)

The ending bill

A Linear head after GAP reads 8 numbers per image. A flatten head on 16 by 16 maps reads 2048.

gap_head = nn.Linear(8, 3)
flat_head = nn.Linear(8 * 16 * 16, 3)
print(sum(p.numel() for p in gap_head.parameters()))
print(sum(p.numel() for p in flat_head.parameters()))

Example

import torch
from torch import nn

torch.manual_seed(0)


def count(model):
    return sum(p.numel() for p in model.parameters())


features = nn.Sequential(
    nn.Conv2d(1, 8, kernel_size=3, padding=1),
    nn.ReLU(),
    nn.Conv2d(8, 8, kernel_size=3, padding=1),
    nn.ReLU(),
)

gap_model = nn.Sequential(
    features,
    nn.AdaptiveAvgPool2d(1),
    nn.Flatten(),
    nn.Linear(8, 3),
)

flat_model = nn.Sequential(
    nn.Sequential(
        nn.Conv2d(1, 8, kernel_size=3, padding=1),
        nn.ReLU(),
        nn.Conv2d(8, 8, kernel_size=3, padding=1),
        nn.ReLU(),
    ),
    nn.Flatten(),
    nn.Linear(8 * 16 * 16, 3),
)

image = torch.rand(2, 1, 16, 16)
maps = features(image)
gap = nn.AdaptiveAvgPool2d(1)(maps)

print("Feature maps:", tuple(maps.shape))
print("After GAP:", tuple(gap.shape))
print("GAP model output:", tuple(gap_model(image).shape))
print("GAP model parameters:", count(gap_model))
print("Flatten model parameters:", count(flat_model))
print("GAP uses fewer:", count(gap_model) < count(flat_model))

Key points

Global average pooling: one mean value for each feature map
Use AdaptiveAvgPool2d(1) to pool any map size down to 1 by 1
Flatten after GAP gives one number per map
A Linear layer after GAP pays weights per map, not per cell
GAP ends a stack with a short summary of each feature map

In this example, the seed is fixed and two images of 16 by 16 run through a small conv stack. The stack leaves eight feature maps of the full image size. AdaptiveAvgPool2d(1) shrinks each map to one cell, and Flatten plus Linear(8, 3) finish the GAP model. Its real parameter count prints as 691. A twin model that flattens the maps instead prints 6811. The GAP ending keeps the same three-class output for far fewer weights.

Note

GAP has no weights of its own to learn

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

import torch
from torch import nn

torch.manual_seed(0)


def count(model):
    return sum(p.numel() for p in model.parameters())


features = nn.Sequential(
    nn.Conv2d(1, 8, kernel_size=3, padding=1),
    nn.ReLU(),
    nn.Conv2d(8, 8, kernel_size=3, padding=1),
    nn.ReLU(),
)

gap_model = nn.Sequential(
    features,
    nn.AdaptiveAvgPool2d(1),
    nn.Flatten(),
    nn.Linear(8, 3),
)

flat_model = nn.Sequential(
    nn.Sequential(
        nn.Conv2d(1, 8, kernel_size=3, padding=1),
        nn.ReLU(),
        nn.Conv2d(8, 8, kernel_size=3, padding=1),
        nn.ReLU(),
    ),
    nn.Flatten(),
    nn.Linear(8 * 16 * 16, 3),
)

image = torch.rand(2, 1, 16, 16)
maps = features(image)
gap = nn.AdaptiveAvgPool2d(1)(maps)

print("Feature maps:", tuple(maps.shape))
print("After GAP:", tuple(gap.shape))
print("GAP model output:", tuple(gap_model(image).shape))
print("GAP model parameters:", count(gap_model))
print("Flatten model parameters:", count(flat_model))
print("GAP uses fewer:", count(gap_model) < count(flat_model))

Exercise

Global average pooling acts on 8 feature maps of 16 by 16. How many numbers describe one image after GAP?

8 128 2048 16