Global Average Pooling in Pytorch

22

torch.mean works on only one dimension.
you needs to use torch.flatten to flatten height and width dimensions before using torch.mean, and use torch.reshape afterward.

23

As of today, the above doesn’t work. nn.AvgPool2d(x) works when x.shape=N * C * H * W.

24

FYI you can do, e.g.:

import torch
feature_maps = torch.rand(16, 512, 7, 7)
feature_vector = feature_maps.mean(dim=(-2, -1))  # or dim=(2, 3)
print(feature_vector.shape)

torch.Size([16, 512])
5 Likes
26

you can use torch.mean() n times for n dimension you need

27

Works like a charm. Thanks!

28

This works but the stride should be kept to default (i.e. the value of the kernel size)

29

Here is a simple example to implement Global Average Pooling:

import torch

import torch.nn as nn

in = torch.randn(10,32,3,3)

pool = nn.AvgPool2d(3)
# note: the kernel size equals the feature map dimensions in the previous layer

output = pool(in)

output = output.squeeze()

print(output.size())
30

nn.functional.avg_pool2d(x, x.size()[2:])

31

PyTorch now has adaptive average pooling:
https://pytorch.org/docs/stable/generated/torch.nn.AdaptiveAvgPool2d.html

https://pytorch.org/docs/stable/generated/torch.nn.functional.adaptive_avg_pool2d.html

1 Like