CODEWITHSUNDEEP
pythonpytorch
Montoya
Montoya

Reputation: 3049

Pytorch - Apply pooling on specific dimension

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?

Upvotes: 4

Views: 2823

Answers (1)

kuzand
kuzand

Reputation: 9806

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)

Upvotes: 6

Related Questions