Pytorch - Apply pooling on specific dimension

Question

I have a 3 dimension vector. I would like to perform a 1d max pool on the second dimension.

According to the documentation of pytorch the pooling is always performed on the last dimension.

https://pytorch.org/docs/stable/nn.html#maxpool1d

For example:

>>> x = torch.rand(5, 64, 32)
>>> pool = nn.MaxPool1d(2, 2)
>>> pool(x).shape
torch.Size([5, 64, 16])

My desired output:

torch.Size([5, 32, 32])

How can i do that?

Answer 1

You can simply permute the dimensions:

x = torch.rand(5, 128, 32)
pool = nn.MaxPool1d(2, 2)
pool(x.permute(0,2,1)).permute(0,2,1)  # shape (5, 128, 32) -> (5, 64, 32)