pytorch api文档:torch.max() vs torch.argmax() 对比

torch.max() vs torch.argmax() 对比

特征 torch.max() torch.argmax()
返回值 无dim:仅最大值;有dim:(最大值, 索引) 仅返回最大值的索引
核心用途 同时需要值和索引时 只需要索引(如分类任务取预测类别)
特殊用法 支持两个张量逐元素取最大值 无
等价关系 torch.max(x, dim=1)[1] = torch.argmax(x, dim=1) 仅能获取索引,无法获取最大值

例子

tensor = torch.tensor([3, 1, 4, 4, 2])

# argmax:只返回最大值索引(第一个最大值的位置)
argmax_idx = torch.argmax(tensor)
# max:返回 (最大值, 最大值索引)
max_val, max_idx = torch.max(tensor)

print(f"argmax索引: {argmax_idx.item()}")  # 输出 2
print(f"max返回值: 数值={max_val.item()}, 索引={max_idx.item()}")  # 输出 4, 2

原文出处:/show_1GW2ercIbwXg.html