我有一个这样的张量：

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>)

我想在两行中获取topk值。目前我能做到的是以下内容：

ipdb>  stacked.topk(2)torch.return_types.topk(values=tensor([[[14.3902, 14.3039]],        [[14.8927, 12.1973]]], grad_fn=<TopkBackward>),indices=tensor([[[60, 12]],        [[12, 23]]]))

从输出中可以看到，从两行中获取了前2个值。我希望得到如下输出：

14.8927 映射到索引 1214.3902 映射到索引 60

请注意，如果前2个值在第一行，它将只返回那里的值并完全忽略第二行，反之亦然。

在这方面需要帮助。

实现我想要的功能的一个非常简陋的方法如下，但这非常简陋，仅展示了BEAM_WIDTH为2的情况：

BEAM_WIDTH = 2top_k = stacked.data.topk(BEAM_WIDTH, dim=2)v1, i1 = top_k[0][0][0], top_k[1][0][0]v2, i2 = top_k[1][0], top_k[1][1][0]i = j = 0final = []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]