max-pooling有很多种实现方式
(1)kmax-pooling
1 import torch2 def kmax_pooling(x, dim, k):3 index = x.topk(k, dim=dim)[1].sort(dim=dim)[0]4 return x.gather(dim, index)5 x = torch.rand(4, 5, 6, 10)6 print(x) # [torch.FloatTensor of size 4x5x6x10]7 y = kmax_pooling(x, 3, 5) 8 print(y) # # [torch.FloatTensor of size 4x5x6x5]