如何查看PyTorch实际计算过程？详解img_t.mean(2)的作用

对于形状为[3,5,5]的张量img_t，PyTorch中mean(dim=2)是沿着第2个维度（索引从0开始）计算算术平均值，具体逻辑如下：

维度对应关系 #

先明确张量各维度的含义（以图像张量为例）：

• 维度0：通道数，共3个通道
• 维度1：图像高度，共5行
• 维度2：图像宽度，共5列

mean(dim=2)的计算逻辑 #

计算时会压缩掉第2个维度，对每个通道的每一行，取该行所有列元素的平均值，最终得到形状为[3,5]的张量：

• 结果张量中mean_2[i,j]的值 = img_t[i,j,0] + img_t[i,j,1] + ... + img_t[i,j,4] 除以5

代码验证示例 #

import torch

# 构造3x5x5的测试张量，数值连续方便计算
img_t = torch.arange(3*5*5).reshape(3,5,5)

# 计算mean(dim=2)
mean_2 = img_t.mean(dim=2)

# 手动验证第0通道第0行的平均值
row_values = img_t[0, 0, :]  # 取值为tensor([0, 1, 2, 3, 4])
manual_avg = row_values.sum() / 5

print(f"原张量形状: {img_t.shape}")
print(f"mean(dim=2)结果形状: {mean_2.shape}")
print(f"手动计算第0通道第0行平均值: {manual_avg}")
print(f"mean(dim=2)对应位置的值: {mean_2[0,0]}")

运行后会看到手动计算结果和mean(dim=2)的输出完全一致，验证了计算逻辑。

和其他维度mean的对比 #

• mean(dim=0)：沿着通道维度（维度0）求平均，压缩维度0，得到[5,5]的张量，每个位置是3个通道对应元素的平均值
• mean(dim=1)：沿着高度维度（维度1）求平均，压缩维度1，得到[3,5]的张量，每个通道的每一列元素的平均值
• mean(dim=2)：沿着宽度维度（维度2）求平均，压缩维度2，得到[3,5]的张量，每个通道的每一行元素的平均值

内容的提问来源于stack exchange，提问作者amnesic