[Python知识库]python代码实现global top-k pooling

python代码实现global top-k pooling：

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import torch
from torch import nn


class TopKPool(nn.Module):
    def __init__(self, k):
        super(TopKPool, self).__init__()
        self.k = k

    def forward(self, x):
        b, c, _, _ = x.shape
        x = x.view(b, c, -1)
        # 在h*w面上取topK大值
        topkv, _ = x.topk(self.k, dim= -1)
        # 对topK个值取平均
        out = topkv.mean(dim= -1)
        return out


if __name__ == "__main__": 
    aa = torch.randn(2, 3, 5, 5)
    print("aa: ", aa.size())
    print("----------------------------")
    print("aa: ", aa)
    # 最大池化
    bb = nn.MaxPool2d(kernel_size=5, stride=2)
    c = bb(aa)
    print("c: ", c)
    print("c.shape", c.shape)

    print("----------------------------")
    # topK-maxpool
    topkmaxpool = TopKPool(k=2)
    dd = topkmaxpool(aa)
    print("dd: ", dd)
    print("dd.shape", dd.shape)