Pytorch - 在特定维度上应用池化

Question

我有一个 3 维向量。我想在第二维上执行一维最大池。

根据 pytorch 的文档，池总是在最后一个维度上执行。

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

例如：

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

我想要的输出：

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

我该怎么做？

Answer 1

您可以简单地置换尺寸：

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)