Additionally, when detecting low-level features like edges, we typically want these features to maintain some degree of translation invariance. For instance, if we have an image X with sharp black-and-white edges and shift the entire image one pixel to the right (Z[i, j] = X[i, j + 1]), the output might differ significantly. In real-world scenarios, objects rarely appear at the exact same pixel locations across different images. Even when photographing a stationery object with a tripod, camera vibrations from shutter movement might shift all objects by one pixel (except with high-end cameras that have special features to prevent this).

Pooling layers address these challenges with a dual purpose: reducing the sensitivity of convolutional layers to position and decreasing the spatial dimensionality of feature representations.

Maximum and Average Pooling

Similar to convolutional layers, pooling operations use a fixed-size window that slides across the input according to a specified stride. For each position the window visits, it computes a single output value. Unlike convolutional layers that perform cross-correlation between input and kernels, pooling layers contain no parameters. The pooling operation is deterministic, typically calculating either the maximum or average of all elements within the window. These operations are known as maximum pooling and average pooling, respectively.

In both cases, the pooling window starts at the top-left corner of the input tensor and moves left to right, top to bottom. At each position, it computes either the maximum or average of the input sub-tensor within the window, depending on whether maximum or average pooling is used.

The output tensor in the example has a height and width of 2, with each element representing the maximum value in its respective pooling window:

Returning to our edge detection example, let's use the output of a convolutional layer as input to a 2×2 max pooling operation. Let X be the convolutional layer input and Y be the pooling layer output. Regardless of whether X[i, j] and X[i, j + 1] have the same value, or whether X[i, j + 1] and X[i, j + 2] have the same value, the pooling layer consistently outputs Y[i, j] = 1. This demonstrates that with max pooling, the convolutional layer can still recognize patterns even when shifted by one element in height or width.

Implementing a Pooling Layer

The following function implements the forward pass for a pooling layer. This is similar to our previous convolution implementation, but without a kernel and by computing either the maximum or average of values in each window.

import torch
from torch import nn
import torch.nn.functional as F

def custom_pool2d(input_tensor, kernel_size, mode='max'):
    """
    Applies 2D pooling operation to the input tensor.
    
    Args:
        input_tensor: Input 2D tensor
        kernel_size: Tuple of (height, width) for pooling window
        mode: Pooling mode ('max' or 'avg')
    
    Returns:
        Pooled tensor
    """
    # Get pooling window dimensions
    window_h, window_w = kernel_size
    # Initialize output tensor
    output_height = input_tensor.shape[0] - window_h + 1
    output_width = input_tensor.shape[1] - window_w + 1
    result = torch.zeros((output_height, output_width))
    
    # Iterate over each position in the output
    for i in range(output_height):
        for j in range(output_width):
            # Extract the current window
            window = input_tensor[i:i+window_h, j:j+window_w]
            # Apply pooling operation
            if mode == 'max':
                result[i, j] = torch.max(window)
            elif mode == 'avg':
                result[i, j] = torch.mean(window)
    
    return result

We can create the input tensor X from the example and verify the output of 2D max pooling:

# Create input tensor
input_data = torch.tensor([[0.0, 1.0, 2.0], 
                          [3.0, 4.0, 5.0], 
                          [6.0, 7.0, 8.0]])

# Apply max pooling
max_pooled = custom_pool2d(input_data, (2, 2))
print(max_pooled)

Output:

tensor([[4., 5.],
        [7., 8.]])

We can also verify average pooling:

# Apply average pooling
avg_pooled = custom_pool2d(input_data, (2, 2), 'avg')
print(avg_pooled)

Output:

tensor([[2., 3.],
        [5., 6.]])

Padding and Stride

Like convolutional layers, pooling layers can modify the output shape. As before, we can use padding and stride to achieve the desired output dimensions. The following example demonstrates padding and stride in a 2D max pooling layer using a built-in function from a deep learning framework. First, we create an input tensor X with four dimensions, where both the batch size and number of channels are 1.

# Create 4D input tensor (batch, channels, height, width)
input_tensor = torch.arange(16, dtype=torch.float32).reshape((1, 1, 4, 4))
print(input_tensor)

Output:

tensor([[[[ 0.,  1.,  2.,  3.],
          [ 4.,  5.,  6.,  7.],
          [ 8.,  9., 10., 11.],
          [12., 13., 14., 15.]]]])

By default, deep learning frameworks use a stride equal to the pooling window size. Therefore, with a pooling window of (3, 3), we get an output with a stride of (3, 3).

# Define 2D max pooling with window size 3
max_pool = nn.MaxPool2d(3)
# Apply pooling to input
pooled_tensor = max_pool(input_tensor)
print(pooled_tensor)

Output:

tensor([[[[10.]]]])

Padding and stride can be manually specified:

# Define pooling with custom padding and stride
max_pool = nn.MaxPool2d(3, padding=1, stride=2)
pooled_tensor = max_pool(input_tensor)
print(pooled_tensor)

Output:

tensor([[[[ 5.,  7.],
          [13., 15.]]]])

We can also use rectangular pooling windows with different padding and stride values for height and width:

# Define pooling with rectangular window and custom parameters
max_pool = nn.MaxPool2d((2, 3), stride=(2, 3), padding=(0, 1))
pooled_tensor = max_pool(input_tensor)
print(pooled_tensor)

Output:

tensor([[[[ 5.,  7.],
          [13., 15.]]]])

Multiple Channels

When processing multi-channel input data, pooling layers operate independently on each input channel, rather than aggregating across channels like convolutional layers. This means the number of output channels in a pooling layer is the same as the number of input channels. Below, we concatenate tensor X and X + 1 along the channel dimension to create a 2-channel input.

# Create 2-channel input
multi_channel_input = torch.cat((input_tensor, input_tensor + 1), 1)
print(multi_channel_input)

Output:

tensor([[[[ 0.,  1.,  2.,  3.],
          [ 4.,  5.,  6.,  7.],
          [ 8.,  9., 10., 11.],
          [12., 13., 14., 15.]],

         [[ 1.,  2.,  3.,  4.],
          [ 5.,  6.,  7.,  8.],
          [ 9., 10., 11., 12.],
          [13., 14., 15., 16.]]]])

As shown below, the number of output channels remains 2 after pooling:

# Apply pooling to multi-channel input
max_pool = nn.MaxPool2d(3, padding=1, stride=2)
pooled_multi = max_pool(multi_channel_input)
print(pooled_multi)

Output:

tensor([[[[ 5.,  7.],
          [13., 15.]],

         [[ 6.,  8.],
          [14., 16.]]]])

Summary

For a given input window, max pooling outputs the maximum value within that window, while average pooling outputs the average value. One of the primary advantages of pooling layers is reducing the position sensitivity of convolutional layers. We can specify padding and stride for pooling layers. Using max pooling with a stride greater than 1 reduces spatial dimensions (height and width). The number of output channels in a pooling layer equals the number of input channels.