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)
三维张量求均值:
import torch
import numpy as np
# ======初始化一个三维矩阵=====
A = torch.ones((4,3,2))
# ======替换三维矩阵里面的值======
A[0] = torch.ones((3,2)) *1
A[1] = torch.ones((3,2)) *2
A[2] = torch.ones((3,2)) *3
A[3] = torch.ones((3,2)) *4
print(A)
B = torch.mean(A ,dim=0)
print(B)
B = torch.mean(A ,dim=1)
print(B)
B = torch.mean(A ,dim=2)
print(B)
输出结果
tensor([[[1., 1.],
[1., 1.],
[1., 1.]],
[[2., 2.],
[2., 2.],
[2., 2.]],
[[3., 3.],
[3., 3.],
[3., 3.]],
[[4., 4.],
[4., 4.],
[4., 4.]]])
tensor([[2.5000, 2.5000],
[2.5000, 2.5000],
[2.5000, 2.5000]])
tensor([[1., 1.],
[2., 2.],
[3., 3.],
[4., 4.]])
tensor([[1., 1., 1.],
[2., 2., 2.],
[3., 3., 3.],
[4., 4., 4.]])