PyTorch torch.nn.AdaptiveAvgPool2d Function

PyTorch torch.nn 参考手册PyTorch torch.nn Reference Manual

torch.nn.AdaptiveAvgPool2dIt is the adaptive average pooling module in PyTorch.

It can pool input feature maps of any size to a specified target size without manually calculating kernel_size and stride.

Function Definition

torch.nn.AdaptiveAvgPool2d(output_size)

Parameter Description:

  • output_size(int or tuple): Output size. Can be (H, W) or a single int (square).

Usage Examples

Example 1: Basic Usage

Pool the feature map to a fixed size:

Example

import torch
import torch.nn as nn

# pool to 1x1
adaptive_pool = nn.AdaptiveAvgPool2d(1)

# input different sizes
x1 = torch.randn(1, 64, 32, 32)
x2 = torch.randn(1, 64, 16, 16)
x3 = torch.randn(1, 64, 8, 8)

out1 = adaptive_pool(x1)
out2 = adaptive_pool(x2)
out3 = adaptive_pool(x3)

print("32x32 ->", out1.shape)
print("16x16 ->", out2.shape)
print("8x8   ->", out3.shape)
print("nAll outputs are pooled to 1x1")

Example 2: Output Arbitrary Size

Pool to non-square:

Example

import torch
import torch.nn as nn

# pool to 4x4
pool = nn.AdaptiveAvgPool2d((4, 4))

x = torch.randn(2, 128, 32, 32)
out = pool(x)

print("Input shape:", x.shape)
print("Output shape:", out.shape)  # (2, 128, 4, 4)

# pool to 1x7 (can be used for sequences)
pool_seq = nn.AdaptiveAvgPool2d((1, 7))
x_seq = torch.randn(2, 64, 10, 20)
out_seq = pool_seq(x_seq)
print("nSequence pooling:", x_seq.shape, "->", out_seq.shape)

Example 3: Use in CNN

Typical global average pooling:

Example

import torch
import torch.nn as nn

class CNNWithGAP(nn.Module):
    def __init__(self, num_classes=10):
        super(CNNWithGAP, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 64, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(64, 128, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(128, 256, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
        )
        # global average pooling
        self.avgpool = nn.AdaptiveAvgPool2d(1)
        self.classifier = nn.Linear(256, num_classes)

    def forward(self, x):
        x = self.features(x)
        x = self.avgpool(x)
        x = x.view(x.size(0), -1)
        x = self.classifier(x)
        return x

model = CNNWithGAP()
x = torch.randn(4, 3, 32, 32)
output = model(x)

print("Input shape:", x.shape)
print("Output shape:", output.shape)
print("nThe feature map is globally pooled into a single vector")

Example 4: Comparison with Regular Pooling

Example

import torch
import torch.nn as nn

# adaptive pooling vs regular pooling
adaptive = nn.AdaptiveAvgPool2d(2)
regular = nn.AvgPool2d(kernel_size=16, stride=16)

x = torch.randn(1, 64, 32, 32)

out_adaptive = adaptive(x)
out_regular = regular(x)

print("Input:", x.shape)
print("Adaptive pooling (output 2x2):", out_adaptive.shape)
print("Regular average pooling (16x16):", out_regular.shape)

Adaptive Pooling vs Regular Pooling

Type Advantages Applicable Scenarios
AdaptiveAvgPool Arbitrary input size Inputs of different scales
AvgPool Faster computation Fixed input size

Frequently Asked Questions

Q1: When to use adaptive pooling?

When the input size is not fixed, or when it needs to be unified to a fixed size.

Q2: What is Global Average Pooling?

AdaptiveAvgPool2d(1), which pools each feature map to a single value.

Use Cases

  • Classification networks: Global average pooling replaces FC layers
  • Multi-scale inputs: Adapt to images of different sizes
  • Feature aggregation: Extract key information

Tip: Global average pooling can significantly reduce the number of parameters.

PyTorch torch.nn 参考手册PyTorch torch.nn Reference Manual

Other Extensions