作用
既能返回最大值，也能返回对应索引
语法
torch.max(input, dim=None, keepdim=False)
参数解释
| 参数 | 作用 | 示例 |
|---|---|---|
input |
输入张量（必填） | 任意维度的张量 |
dim |
指定计算维度（可选） | 详见下面 |
keepdim |
是否保留原维度（默认False） | keepdim=True 时，结果维度和输入一致 |
dim参数解释：
- 不指定维度：返回张量中 全局最大值（标量）；
- 指定维度：
dim=0按列，dim=1按行。返回一个元组(最大值, 最大值对应的索引)，其中：- 第一个元素是指定维度上的最大值张量；
- 第二个元素是这些最大值对应的索引（等价于同维度下的
torch.argmax()结果）。
例子
import torch
# 1. 一维张量（向量）
tensor_1d = torch.tensor([3, 1, 4, 2, 4])
# 不指定dim：只返回全局最大值
max_val = torch.max(tensor_1d)
# 指定dim=0（一维张量只有dim=0）：返回 (最大值, 索引)
max_val_idx = torch.max(tensor_1d, dim=0)
print(f"一维张量: {tensor_1d}")
print(f"全局最大值（无dim）: {max_val.item()}") # 输出 4
print(f"指定dim=0的结果: 最大值={max_val_idx[0].item()}, 索引={max_val_idx[1].item()}") # 索引=2（第一个4的位置）
# 2. 二维张量（矩阵）
tensor_2d = torch.tensor([
[1, 5, 3],
[2, 4, 6]
])
# 2.1 不指定dim：返回全局最大值
max_2d_global = torch.max(tensor_2d)
print(f"\n二维张量:\n{tensor_2d}")
print(f"全局最大值（无dim）: {max_2d_global.item()}") # 输出 6
# 2.2 指定dim=0（按列计算）：每列的最大值 + 对应行索引
max_2d_dim0 = torch.max(tensor_2d, dim=0)
print(f"\ndim=0（按列）的最大值: {max_2d_dim0[0]}") # 输出 [2,5,6]
print(f"dim=0（按列）的最大值索引: {max_2d_dim0[1]}") # 输出 [1,0,1]（和argmax(dim=0)结果一致）
# 2.3 指定dim=1（按行计算）：每行的最大值 + 对应列索引
max_2d_dim1 = torch.max(tensor_2d, dim=1)
print(f"\ndim=1（按行）的最大值: {max_2d_dim1[0]}") # 输出 [5,6]
print(f"dim=1（按行）的最大值索引: {max_2d_dim1[1]}") # 输出 [1,2]（和argmax(dim=1)结果一致）
# 2.4 keepdim=True：保留原维度（结果形状和输入匹配）
max_2d_dim1_keep = torch.max(tensor_2d, dim=1, keepdim=True)
print(f"\ndim=1且keepdim=True:")
print(f"最大值: \n{max_2d_dim1_keep[0]}") # 输出 [[5], [6]]
print(f"索引: \n{max_2d_dim1_keep[1]}") # 输出 [[1], [2]]
输出结果：
一维张量: tensor([3, 1, 4, 2, 4])
全局最大值（无dim）: 4
指定dim=0的结果: 最大值=4, 索引=2
二维张量:
tensor([[1, 5, 3],
[2, 4, 6]])
全局最大值（无dim）: 6
dim=0（按列）的最大值: tensor([2, 5, 6])
dim=0（按列）的最大值索引: tensor([1, 0, 1])
dim=1（按行）的最大值: tensor([5, 6])
dim=1（按行）的最大值索引: tensor([1, 2])
dim=1且keepdim=True:
最大值:
tensor([[5],
[6]])
索引:
tensor([[1],
[2]])
实战场景
torch.max() 在深度学习中最常用的场景是分类任务，同时它还能实现“两个张量逐元素取最大值”的特殊用法：
场景1：分类任务中获取预测值和类别
# 模拟模型输出（batch_size=3，类别数=4）
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得分最高
])
# 按行取最大值（得分）和索引（类别）
pred_scores, pred_classes = torch.max(model_output, dim=1)
print(f"模型输出得分:\n{model_output}")
print(f"\n每个样本的预测得分: {pred_scores}") # 输出 [0.8000, 0.9000, 0.7000]
print(f"每个样本的预测类别索引: {pred_classes}") # 输出 [1,2,0]
场景2：两个张量逐元素取最大值（特殊用法）
torch.max() 还有一个重载形式：torch.max(input, other)，用于逐元素比较两个张量，返回更大的值：
a = torch.tensor([1, 5, 3])
b = torch.tensor([2, 4, 6])
# 逐元素取最大值
max_ab = torch.max(a, b)
print(f"\n张量a: {a}")
print(f"张量b: {b}")
print(f"逐元素最大值: {max_ab}") # 输出 [2,5,6]
与 torch.argmax() 区别
详见链接
总结
torch.max()是多功能工具：无维度时返回全局最大值，指定维度时返回(最大值, 索引)元组，还能逐元素比较两个张量；- 分类任务中，
torch.max(output, dim=1)[1]等价于torch.argmax(output, dim=1)，前者能同时拿到预测得分，后者更简洁； - 记住核心区别：
max能拿“值+索引”，argmax只拿“索引”，按需选择即可。