nn.MaxPool2d and nn.AvgPool2d in PyTorch

Introduction

Pooling layers help reduce the size of images or feature maps in neural networks. MaxPool2d picks the biggest value in a small area, while AvgPool2d takes the average. This makes the model faster and focuses on important features.

When you want to reduce the size of image data to make the model faster.

When you want to keep the strongest features in an image using max pooling.

When you want to smooth features by averaging nearby values using average pooling.

When building convolutional neural networks for image recognition.

When you want to reduce noise in feature maps by averaging.

Syntax

PyTorch

nn.MaxPool2d(kernel_size, stride=None, padding=0, dilation=1, return_indices=False, ceil_mode=False)
nn.AvgPool2d(kernel_size, stride=None, padding=0, ceil_mode=False, count_include_pad=True, divisor_override=None)

kernel_size is the size of the window to pool over (e.g., 2 means 2x2 window).

stride is how far the window moves each step. If None, it equals kernel_size.

Examples

Max pooling with a 2x2 window and stride 2 (default). It picks the max value in each 2x2 block.

PyTorch

nn.MaxPool2d(2)

Average pooling with a 3x3 window, moving 1 step each time, and padding 1 to keep size.

PyTorch

nn.AvgPool2d(kernel_size=3, stride=1, padding=1)

Max pooling that also returns the indices of max values, useful for unpooling later.

PyTorch

nn.MaxPool2d(kernel_size=2, stride=2, return_indices=True)

Sample Model

This code shows how max pooling picks the biggest number in each 2x2 block, and average pooling calculates the average of each 2x2 block.

PyTorch

import torch
import torch.nn as nn

# Create a sample input tensor (1 image, 1 channel, 4x4 size)
input_tensor = torch.tensor([[[[1.0, 2.0, 3.0, 4.0],
                               [5.0, 6.0, 7.0, 8.0],
                               [9.0, 10.0, 11.0, 12.0],
                               [13.0, 14.0, 15.0, 16.0]]]])

# Define MaxPool2d with 2x2 kernel and stride 2
max_pool = nn.MaxPool2d(kernel_size=2, stride=2)
# Define AvgPool2d with 2x2 kernel and stride 2
avg_pool = nn.AvgPool2d(kernel_size=2, stride=2)

# Apply max pooling
max_pooled = max_pool(input_tensor)
# Apply average pooling
avg_pooled = avg_pool(input_tensor)

print("Input Tensor:")
print(input_tensor)
print("\nMax Pooled Output:")
print(max_pooled)
print("\nAverage Pooled Output:")
print(avg_pooled)

Important Notes

MaxPool2d helps keep the strongest features by picking the highest value.

AvgPool2d smooths the features by averaging, which can reduce noise.

Pooling reduces the size of data, making models faster and less likely to overfit.

Summary

MaxPool2d picks the maximum value in each window to keep strong signals.

AvgPool2d calculates the average value in each window to smooth features.

Both reduce data size and help neural networks focus on important information.