torch.mean()函数

torch.mean(input) 输出input 各个元素的的均值

1. 不指定任何参数就是所有元素的算术平均值
2. 指定参数可以计算每一行或者 每一列的算术平均数

对于1 如：

import torch
a=torch.randn(3)  #随机生成一个一维的矩阵
b=torch.randn(1,3)  #随机生成一个二维的矩阵
print(a)
print(b)
torch.mean(a)

结果

tensor([ 0.1613,  0.7873, -0.3453])

tensor([[ 0.5759,  1.6183, -0.2187]])

tensor(0.2011)

对于2 如果需要指定参数的话，需要使用dim 如：

import torch
a=torch.randn(4,4)
print(a)
b=torch.mean(a,dim=0,keepdim=True)
print(b)
c=torch.mean(a,dim=1,keepdim=True)
print(c)

结果为：

tensor([[-0.3011, -0.8956,  1.4941, -0.1555],
        [-2.4680, -0.9700,  0.1180,  0.1979],
        [ 1.0346, -0.2269,  1.7756,  0.6333],
        [-0.0050,  0.4847, -0.3076,  0.4797]])
tensor([[-0.4349, -0.4019,  0.7700,  0.2888]])
tensor([[ 0.0355],
        [-0.7805],
        [ 0.8041],
        [ 0.1629]])

dim
指定为1时，求得是行的平均值；
指定为0时，求得是列的平均值。