多维张量的前 K 个索引

我有一个 2D 张量，我想获取前 k 个值的索引。我知道关于pytorch 的 topk功能。 pytorch 的 topk 函数的问题是，它计算某个维度上的 topk 值。我想获得两个维度上的 topk 值。

例如对于以下张量

a = torch.tensor([[4, 9, 7, 4, 0],
        [8, 1, 3, 1, 0],
        [9, 8, 4, 4, 8],
        [0, 9, 4, 7, 8],
        [8, 8, 0, 1, 4]])

pytorch 的 topk 函数会给我以下内容。

values, indices = torch.topk(a, 3)

print(indices)
# tensor([[1, 2, 0],
#        [0, 2, 1],
#        [0, 1, 4],
#        [1, 4, 3],
#        [1, 0, 4]])

但我想得到以下内容

tensor([[0, 1],
        [2, 0],
        [3, 1]])

这是二维张量中 9 的索引。

有没有什么方法可以使用 pytorch 来实现这一点？

v, i = torch.topk(a.flatten(), 3)
print (np.array(np.unravel_index(i.numpy(), a.shape)).T)

Output:

[[3 1]
 [2 0]
 [0 1]]

1. 压平并找到top k
2. 使用以下命令将 1D 索引转换为 2Dunravel_index