PyTorch 实现kmax-pooling

max-pooling有很多种实现方式

(1)kmax-pooling

1 import torch
2 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]

 

    原文作者:pytorch
    原文地址: https://www.cnblogs.com/Joyce-song94/p/7277871.html
    本文转自网络文章,转载此文章仅为分享知识,如有侵权,请联系博主进行删除。
点赞