torch.mean是对torch.Tensor取均值。具体怎么使用，请看下面的例子

不指定维度求均值

如果不指定维度，是对Tensor里所有元素求均值，最后计算出来的结果是一个标量。

a = torch.Tensor([0, 1, 2, 3, 4, 5]).view(2, 3, 1)
print(a)

mean = torch.mean(a)
print(mean, mean.shape)

结果如下：

指定其中一维

如果指定了某一维度，那么剩下的所有元素算均值。
下面的例子是对（2，3，1）Tensor的0维做均值，N=2,即所选的dim，输出为剩下的维度（3，1）
第一个：（0+3)/2=1.5
第二个： (1+4)/2=2.5
第三个： (2+5)/2=3.5

a = torch.Tensor([0, 1, 2, 3, 4, 5]).view(2, 3, 1)
print(a)

mean = torch.mean(a, 0)
print(mean, mean.shape)

例子2

如下对dim 1做mean

N=3,即所选的dim，输出为剩下的维度（2，1）
第一个：（0+1+2)/3=1
第二个： (3+4+5)/3=4

a = torch.Tensor([0, 1, 2, 3, 4, 5]).view(2, 3, 1)
print(a)

mean = torch.mean(a, 1)
print(mean, mean.shape)

例子3

如下对dim 2做mean

N=1,即所选的dim，输出为剩下的维度（3，2）
[0][0]：（0)/1=0
[0][1]：（1)/1=1
[1][0]：（2)/1=2
[1][1]：（3)/1=3
[2][0]：（4)/1=4
[2][1]：（5)/1=5

a = torch.Tensor([0, 1, 2, 3, 4, 5]).view(2, 3, 1)
print(a)

mean = torch.mean(a, 2)
print(mean, mean.shape)

指定多维

其实跟指定一个维度一样，下面如例子

dim=（1，2）
那么N = 2*3 = 6, 输出为（1，）

[0]=(0+1+2+3+4+5)/6 = 2.5

a = torch.Tensor([0, 1, 2, 3, 4, 5]).view(2, 3, 1)
print(a)

mean = torch.mean(a, (0,1))
print(mean, mean.shape)

如果需要保持维度，这里的输入为三个维度，那么保持维度的意思就是输出也是三个维度。
那么可以加入keepdim=True，
如下：

a = torch.Tensor([0, 1, 2, 3, 4, 5]).view(2, 3, 1)
print(a)

mean = torch.mean(a, (0,1), keepdim=True)
print(mean, mean.shape)