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