Implement Global Average Pooling

Problem

Implement Global Average Pooling. Given a feature map tensor of shape (batch_size, channels, height, width), compute the spatial average for each channel, reducing the spatial dimensions to 1x1. This replaces fully connected layers at the end of CNNs.

Solution

import numpy as np

def global_average_pooling(x: np.ndarray) -> np.ndarray:
    # x shape: (batch_size, channels, height, width)
    # Output shape: (batch_size, channels)
    return np.mean(x, axis=(2, 3))

Explanation

1. Input: A 4D tensor (B, C, H, W) representing a batch of feature maps from a convolutional layer.
2. Operation: Average all spatial positions (H, W) for each channel independently, using np.mean over axes 2 and 3.
3. Output: A 2D tensor (B, C) where each value is the average of an entire feature map.
4. Benefits: Reduces parameters drastically compared to flattening + fully connected layers, acts as a structural regularizer, and makes the network invariant to input spatial dimensions.

Complexity

• Time: O(B C H * W) to compute the averages
• Space: O(B * C) for the output