介绍
PyTorch 中张量（Tensor）的 mean() 方法，是张量操作中非常常用的求 平均值函数。
作用
计算张量中元素的 平均值，可以指定维度（dim）来计算指定维度的均值，也可以直接计算整个张量的均值。
用法
mean(dim)
参数解释
| 参数 | 作用 |
|---|---|
dim |
可选，沿着指定维度（整数/元组）计算均值，不指定则计算整个张量的均值。-1 表示沿着最后一个维度（倒数第一个维度）计算均值 |
keepdim |
可选，默认 False，若为 True，计算后保留原维度（指定维度变为1） |
dtype |
可选，指定输出均值的数据类型（如 torch.float64） |
例子
import torch
# 创建示例张量
tensor = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
1. 计算整个张量的均值（无参数）
mean_all = tensor.mean()
print("整个张量的均值:", mean_all) # 输出: tensor(3.5)
2. 按指定维度计算均值（dim参数）
# dim=0: 按列计算均值（纵向）
mean_dim0 = tensor.mean(dim=0)
print("按列（dim=0）的均值:", mean_dim0) # 输出: tensor([2.5000, 3.5000, 4.5000])
# dim=1: 按行计算均值（横向）
mean_dim1 = tensor.mean(dim=1)
print("按行（dim=1）的均值:", mean_dim1) # 输出: tensor([2., 5.])
3. 保留维度（keepdim=True）：结果维度与原张量一致（仅指定维度变为1）
mean_dim1_keepdim = tensor.mean(dim=1, keepdim=True)
print("保留维度的行均值:", mean_dim1_keepdim)
# 输出: tensor([[2.],
# [5.]])
注意
数据类型限制
mean() 仅支持浮点型张量（float32/float64 等），整数型张量直接调用会报错，需要先转换类型：
# 错误示例：整数张量直接求均值
int_tensor = torch.tensor([1, 2, 3])
# int_tensor.mean() # 会报错：RuntimeError: mean is not implemented for type torch.int64
# 正确做法：转换为浮点型
float_tensor = int_tensor.float()
print(float_tensor.mean()) # 输出: tensor(2.)
多维张量的维度指定
对于高维张量（如 3 维），可以指定多个维度同时计算均值：
# 3维张量：shape=(2, 3, 4)
tensor_3d = torch.randn(2, 3, 4)
# 计算dim=1和dim=2的均值，结果shape=(2,)
mean_2d = tensor_3d.mean(dim=(1, 2))
print("3维张量指定多维度均值的shape:", mean_2d.shape) # 输出: torch.Size([2])
忽略 NaN 值
如果张量中有 NaN，普通 mean() 会返回 NaN，需要用 nanmean() 替代：
nan_tensor = torch.tensor([1.0, 2.0, float('nan'), 4.0])
print("普通mean（含NaN）:", nan_tensor.mean()) # 输出: tensor(nan)
print("nanmean（忽略NaN）:", torch.nanmean(nan_tensor)) # 输出: tensor(2.3333)
总结
tensor.mean()是 PyTorch 计算张量均值的核心方法，无参数时计算全局均值，指定dim可计算指定维度的均值；- 注意数据类型：仅支持浮点型张量，整数张量需先转换为
float； - 特殊场景：
keepdim=True保留维度，nanmean()忽略NaN值，多维度可通过dim=(d1, d2)指定。