Pytorch torch.mean()的简单用法

简单来说就是求平均数。
比如以下的三种简单情况：

import torch

x1 = torch.Tensor([1, 2, 3, 4])
x2 = torch.Tensor([[1],
                   [2],
                   [3],
                   [4]])
x3 = torch.Tensor([[1, 2],
                   [3, 4]])
y1 = torch.mean(x1)
y2 = torch.mean(x2)
y3 = torch.mean(x3)
print(y1)
print(y2)
print(y3)

输出：

tensor(2.5000)
tensor(2.5000)
tensor(2.5000)

也就是说，在没有指定维度的情况下，就是对所有数进行求平均。

更多的时候用到的是有维度的情形，如：

import torch

x = torch.Tensor([1, 2, 3, 4, 5, 6]).view(2, 3)
y_0 = torch.mean(x, dim=0)
y_1 = torch.mean(x, dim=1)
print(x)
print(y_0)
print(y_1)

输出：

tensor([[1., 2., 3.],
        [4., 5., 6.]])
tensor([2.5000, 3.5000, 4.5000])
tensor([2., 5.])

输入tensor的形状为(2, 3)，其中2为第0维，3为第1维。对第0维求平均，得到的结果为形状为(1, 3)的tensor；对第1维求平均，得到的结果为形状为(2, 1)的tensor。
可以理解为，对哪一维做平均，就是将该维所有的数做平均，压扁成1层(实际上这一层就给合并掉了，比如上面的例子，2维的tensor在求平均数后变成了1维)，而其他维的形状不影响。
如果要保持维度不变(例如在深度网络中)，则可以加上参数keepdim=True：

y = torch.mean(x, dim=1, keepdim=True)