Pytorch - применить пул по заданному измерению c

Question

У меня есть 3-мерный вектор. Я хотел бы выполнить 1d max pool для второго измерения.

В соответствии с документацией 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)