PyTorch 实现kmax-pooling

kmax-pooling 是一种池化方式，它可以保留每个通道的前 k 个最大值，这个方法在处理文本数据时特别有用。下面是 PyTorch 实现 kmax-pooling 的代码：

import torch

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

    def forward(self, x):
        # 将输入的张量沿着通道维度排序
        sorted_tensor, indices = torch.sort(x, dim=-1, descending=True)
        # 取出每个通道的前 k 个最大值
        top_k = sorted_tensor[:, :, :self.k]
        # 对这些最大值进行排序并获取它们在原始张量中的索引
        top_k_sorted, _ = torch.sort(top_k, dim=-1)
        indices = indices.gather(dim=-1, index=torch.argsort(top_k, dim=-1))
        # 将索引扩展到与 top_k_sorted 一样的形状，并使用这些索引从原始张量中取出最大值
        max_k = torch.gather(x, dim=-1, index=indices)
        return max_k

在 forward 函数中，我们首先对输入张量 x 沿着通道维度进行排序，然后取出每个通道的前 k 个最大值。我们还需要对这些最大值进行排序并获取它们在原始张量中的索引。最后，我们将这些索引扩展到与 top_k_sorted 一样的形状，并使用这些索引从原始张量中取出最大值。