N 个最高值的 Torch argmax [重复]

Question

我想做一些类似 argmax 但有多个最高值的事情。我知道如何使用普通的 torch.argmax

>>> a = torch.randn(4, 4)
>>> a
tensor([[ 1.3398,  1.2663, -0.2686,  0.2450],
        [-0.7401, -0.8805, -0.3402, -1.1936],
        [ 0.4907, -1.3948, -1.0691, -0.3132],
        [-1.6092,  0.5419, -0.2993,  0.3195]])
>>> torch.argmax(a)
tensor(0)

但现在我需要找到前 N 个值的索引。所以像这样的

>>> a = torch.randn(4, 4)
>>> a
tensor([[ 1.3398,  1.2663, -0.2686,  0.2450],
        [-0.7401, -0.8805, -0.3402, -1.1936],
        [ 0.4907, -1.3948, -1.0691, -0.3132],
        [-1.6092,  0.5419, -0.2993,  0.3195]])
>>> torch.argmax(a,top_n=2)
tensor([0,1])

我在 pytorch 中没有找到任何能够做到这一点的函数，有人知道吗？

Answer 1

太棒了！所以你需要张量的前 k 个最大元素。

[答案 1] 无论维度如何，您都需要所有元素中最大的前 k 个。因此，展平张量并使用 torch.topk 函数获取前 3 个（例如）元素的索引：

>>> a = torch.randn(5,4)
>>> a
tensor([[ 0.8292, -0.5123, -0.0741, -0.3043],
        [-0.4340, -0.7763,  1.9716, -0.5620],
        [ 0.1582, -1.2000,  1.0202, -1.5202],
        [-0.3617, -0.2479,  0.6204,  0.2575],
        [ 1.8025,  1.9864, -0.8013, -0.7508]])
>>> torch.topk(a.flatten(), 3).indices
tensor([17,  6, 16])

[答案 2] 您需要给定输入张量沿给定维度的 k 个最大元素。因此，请参阅函数 torch.topk 给定 here 的 PyTorch 文档。

Answer 2

相信你在找torch.topk