Python 如何在PyTorch中跨多行检索topk值及其各自的索引？

我有一个像这样的张量：

tensor([[[ 7.3478, -1.8058, -2.6140,  ..., -0.2719, -0.3171, -0.4737]],

    [[ 7.3606, -1.8269, -1.9825,  ..., -0.8680,  0.4894,  0.2708]]],
   grad_fn=<CatBackward>)

从输出中，您可以看到前2个值是从两行中检索到的。我希望得到如下输出：

14.8927 that maps to index 12
14.3902 that maps to index 60

请注意，如果前2个值位于第一行中，它将只返回其中的值，并完全忽略第二行，反之亦然

在这方面需要帮助

我想说的是一种非常老套的方法，如下所示，但它非常老套，并且在梁的宽度为2时显示：

BEAM_WIDTH = 2
top_k = stacked.data.topk(BEAM_WIDTH, dim=2)

v1, i1 = top_k[0][0][0], top_k[1][0][0]
v2, i2 = top_k[0][1][0], top_k[1][1][0]

i = j = 0
final = []
for _ in range(BEAM_WIDTH):
    if v1[i] >= v2[j]:
        final.append((v1[i], i1[i]))
        i += 1
    else:
        final.append((v2[j], i2[j]))
        j += 1

---

重复索引

我相信这就是你想要的。首先，在展开列表中找到topk元素，然后将这些索引转换回行相对格式

topk_values, linear_indices = stacked.flatten().topk(2)
topk_indices = linear_indices % stacked.shape[-1]

---

唯一索引

前面的方法不强制执行唯一索引。如果需要唯一的索引，那么您可以找到行之间的最大值，然后找到其中的topk

topk_values, topk_indices = stacked.max(dim=0)[0].flatten().topk(2)

---

示例

为了演示这两种方法之间的差异，假设您

stacked = torch.tensor([[[11,8,0]],
                        [[10,9,0]]])

在重复索引的情况下，您将以

topk_values=[11, 10]
topk_indices=[0, 0]

在唯一索引的情况下，您将得到

topk_values=[11, 9]
topk_indices=[0, 1]

---

@jodag我刚刚更新了问题，以便更好地解释我想做什么。希望这有助于澄清问题。如果前两个值位于不同行的同一列中，会发生什么情况？@jodag在这种情况下，它应该返回不同行的值，索引应该与相应的值匹配。它不应该是

topk\u index=linear\u index%叠加。shape[-1]

而不是

x

？是的，现在修好了

topk_values=[11, 9]
topk_indices=[0, 1]