PyTorch如何判断张量A前k大元素索引匹配张量B的argmax索引

错误原因 #

你的代码问题出在torch.any(top_idx, 1)的使用逻辑完全错误：

• torch.any()作用是对布尔张量沿指定维度判断是否存在True，当输入为整数张量时，会自动将所有非0值识别为True、0值识别为False，返回的是对应维度下是否有非0值的布尔结果，根本不是你想要的「判断真实索引是否在topk索引列表中」的逻辑。
• 以示例第一行为例：top_idx是[2, 1]，torch.any(top_idx, 1)返回True（对应整数值1），而该行真实索引是2，1 == 2不成立，所以第一行被误判为不匹配，最终只返回第二行的索引。

正确实现 #

思路是先将真实索引扩维后和topk索引批量比较，再逐行判断是否存在匹配：

import torch

# 示例输入
A = torch.tensor([[0.1, 0.4, 0.5],
                  [0.5, 0.4, 0.1],
                  [0.4, 0.5, 0.1]])
B = torch.tensor([[0, 0, 1],
                  [0, 1, 0],
                  [0, 0, 1]])
k = 2

# 1. 得到每行真实类别索引
y_true = B.argmax(dim=1)
# 2. 得到每行topk最大值的索引
top_idx = torch.topk(A, k, dim=1).indices
# 3. 扩维后批量比较，逐行判断是否存在匹配
is_match = (top_idx == y_true.unsqueeze(1)).any(dim=1)
# 4. 提取匹配的行索引
result = torch.where(is_match)[0]

print(result) # 输出 tensor([0, 1])

内容的提问来源于stack exchange，提问作者Scoops