torch.topk与torch.tensor.max()疑问：最大值匹配但索引不一致

问题现象 #

对比torch.topk与torch.tensor.max()的输出时，发现两个规律：

• 所有情况下，两者返回的最大值完全相等
• 最大值对应的索引却偶现不匹配

用torch.equal验证：

• 最大值（max()返回的x与topk()返回的y[0,:,:]）完全匹配
• 索引（max()返回的prediction与topk()返回的top_indices[0,:,:]）偶现不匹配

参考代码 #

for step,batch_data in enumerate(val_loader):
    inputs = batch_data[0].to("cuda")
    target = batch_data[1].squeeze().numpy()

    with torch.no_grad():
        outputs = model(inputs)

    x,prediction = outputs.squeeze().max(0)
    y,top_indices = torch.topk(outputs.squeeze(),k=2,dim=0)

问题复现代码（偶现不匹配） #

for i in range(100):
    demo_tensor = torch.randn(3,256,256).to("cuda")
    max_1,coord_1 = demo_tensor.max(dim=0)
    max_2,coord_2 = torch.topk(demo_tensor,k=2,dim=0)
    if torch.equal(max_1,max_2[0,:,:]) and torch.equal(coord_1,coord_2[0,:,:]):
        print("Values match")
    else:
        print("No match")

    print("\n")

原因分析 #

这个问题的核心是当张量中存在多个相同的最大值时，两个函数的索引处理逻辑不同：

• max(dim=0)会返回该维度上第一个出现的最大值的索引
• topk(k=2, dim=0)在对相同值排序时，由于CUDA并行计算的特性，相同值的索引顺序可能存在不确定性，不会严格遵循"第一个出现"的规则，因此偶现索引不匹配的情况。

可以构造一个含重复最大值的张量验证：

demo_tensor = torch.tensor([[5, 3], [5, 4], [2, 5]]).to("cuda")
max_1, coord_1 = demo_tensor.max(dim=0)
max_2, coord_2 = torch.topk(demo_tensor, k=2, dim=0)
print("max索引:", coord_1)
print("topk第一个索引:", coord_2[0])

运行后会看到两者索引的差异。

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