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.
As of today, the above doesn’t work. nn.AvgPool2d(x) works when x.shape=N * C * H * W.
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])
you can use torch.mean() n times for n dimension you need
Works like a charm. Thanks!
This works but the stride should be kept to default (i.e. the value of the kernel size)
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())
nn.functional.avg_pool2d(x, x.size()[2:])
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