作用
返回张量中指定维度上最大值所在的索引。
应用场景
分类任务中，模型输出的是各类别的概率/得分，用 argmax 就能找到“预测类别”（得分最高的类别索引）；
语法
torch.argmax(input, dim=None, keepdim=False)
参数解释
| 参数 | 作用 | 示例 |
|---|---|---|
input |
输入张量（必填） | 任意维度的张量 |
dim |
指定计算维度（可选） | dim=0 按列，dim=1 按行，不指定则展平成一维，再返回全局最大值的索引 |
keepdim |
是否保留原维度（默认False） | keepdim=True 时，结果维度和输入一致 |
注意
argmax返回的是索引（整数），不是最大值本身；如果要同时获取最大值和索引，用torch.max()；- 如果有多个相同的最大值，
argmax只会返回第一个最大值的索引； - 索引从 0 开始（和Python列表索引规则一致）。
例子
import torch
# 1. 一维张量（向量）
tensor_1d = torch.tensor([3, 1, 4, 2])
argmax_1d = torch.argmax(tensor_1d)
print(f"一维张量: {tensor_1d}")
print(f"全局最大值索引: {argmax_1d.item()}") # 输出 2（值为4的位置）
# 2. 二维张量（矩阵）
tensor_2d = torch.tensor([
[1, 5, 3],
[2, 4, 6]
])
# 2.1 不指定dim：展平后找全局最大值索引
argmax_2d_global = torch.argmax(tensor_2d)
print(f"\n二维张量:\n{tensor_2d}")
print(f"全局最大值索引（展平后）: {argmax_2d_global.item()}") # 输出 5（展平后[1,5,3,2,4,6]中6的位置）
# 2.2 指定dim=0：按列找最大值索引（每列的最大值在第几行）
argmax_2d_dim0 = torch.argmax(tensor_2d, dim=0)
print(f"dim=0（按列）的最大值索引: {argmax_2d_dim0}") # 输出 [1,0,1]
# 2.3 指定dim=1：按行找最大值索引（每行的最大值在第几列）
argmax_2d_dim1 = torch.argmax(tensor_2d, dim=1)
print(f"dim=1（按行）的最大值索引: {argmax_2d_dim1}") # 输出 [1,2]
# 3. keepdim=True：保留原维度（结果形状和输入匹配）
argmax_2d_dim1_keep = torch.argmax(tensor_2d, dim=1, keepdim=True)
print(f"\ndim=1且keepdim=True: \n{argmax_2d_dim1_keep}") # 输出 [[1], [2]]（保持2行1列）
输出结果：
一维张量: tensor([3, 1, 4, 2])
全局最大值索引: 2
二维张量:
tensor([[1, 5, 3],
[2, 4, 6]])
全局最大值索引（展平后）: 5
dim=0（按列）的最大值索引: tensor([1, 0, 1])
dim=1（按行）的最大值索引: tensor([1, 2])
dim=1且keepdim=True:
tensor([[1],
[2]])
实战场景
这是 torch.argmax 最核心的使用场景——分类任务中的应用，从模型输出中提取预测类别：
# 模拟模型输出（batch_size=3，类别数=4）
# 每一行代表1个样本，每一列代表对应类别的得分
model_output = torch.tensor([
[0.1, 0.8, 0.05, 0.05], # 样本1：类别1得分最高
[0.0, 0.1, 0.9, 0.0], # 样本2：类别2得分最高
[0.7, 0.2, 0.05, 0.05] # 样本3：类别0得分最高
])
# 按行（dim=1）找最大值索引 → 每个样本的预测类别
pred_classes = torch.argmax(model_output, dim=1)
print(f"模型输出得分:\n{model_output}")
print(f"\n每个样本的预测类别索引: {pred_classes}") # 输出 [1,2,0]
# 结合类别标签映射（比如标签：0=猫,1=狗,2=鸟,3=鱼）
label_map = {0: '猫', 1: '狗', 2: '鸟', 3: '鱼'}
pred_labels = [label_map[idx.item()] for idx in pred_classes]
print(f"预测类别名称: {pred_labels}") # 输出 ['狗', '鸟', '猫']
与 torch.max() 区别
详见链接
总结
torch.argmax()核心是返回张量指定维度上最大值的索引，不指定维度则展平后找全局索引；- 分类任务中，常用
torch.argmax(output, dim=1)从模型输出（batch×类别数）中提取每个样本的预测类别； - 区分
argmax（只返回索引）和max（返回值+索引），按需选择即可。