ML Katas

Global Average Pooling

easy (<10 mins) cnn einops reduction
this year by E

Description

Implement global average pooling, a common operation in convolutional neural networks. For a batch of feature maps of shape (B, C, H, W), you need to compute the mean of each channel, resulting in a tensor of shape (B, C).

Starter Code

import torch
from einops import reduce

def global_average_pool(feature_maps):
    # Your einops code here
    pass

Verification

Create a tensor of shape (16, 512, 7, 7). The output should have the shape (16, 512).