利用pytorch，将输入特征[m,n]，每行topk置1，其余值置零

下面是一个示例代码，其中输入特征为x，topk的数量为k：

import torch

def topk_mask(x, k):
    # 对每行进行topk操作
    values, indices = torch.topk(x, k, dim=1)

    # 创建一个与x形状相同的全零张量
    mask = torch.zeros_like(x)

    # 将每行topk位置设为1
    mask.scatter_(1, indices, 1)

    return mask

使用示例：

# 创建一个4行5列的随机张量
x = torch.randn(4, 5)

# 对每行取最大的2个值，其余位置设为0
mask = topk_mask(x, 2)

print(x)
print(mask)

输出：

tensor([[-1.3807, -0.1081, -0.2223, -0.1873,  0.3692],
        [ 0.4514, -0.2676,  0.4349,  0.4448, -0.9266],
        [ 0.5778, -0.1975, -0.1772, -0.6595,  1.2591],
        [ 1.3213, -0.5793,  0.0719, -0.3505, -0.3875]])
tensor([[0., 0., 0., 0., 1.],
        [1., 0., 1., 1., 0.],
        [1., 0., 0., 0., 1.],
        [1., 0., 0., 0., 0.]])

可以看到，每行最大的2个值被置为1，其余位置为0。