PyTorch 实现kmax-pooling

max-pooling有很多种实现方式

（1）kmax-pooling

 import torch
 def kmax_pooling(x, dim, k):
     index = x.topk(k, dim=dim)[1].sort(dim=dim)[0]
     return x.gather(dim, index)
 x = torch.rand(4, 5, 6, 10)
 print(x)                        # [torch.FloatTensor of size 4x5x6x10]
 y = kmax_pooling(x, 3, 5)
 print(y)                        # # [torch.FloatTensor of size 4x5x6x5]