pytorch中tensor.mean(axis, keepdim)

import numpy as np
import torch

x=[
[[1,2,3,4],
 [5,6,7,8],
 [9,10,11,12]],

[[13,14,15,16],
 [17,18,19,20],
 [21,22,23,24]]
]
x=torch.tensor(x).float()
#
print("shape of x:")  ##[2,3,4]
print(x.shape)
#
print("shape of x.mean(axis=0,keepdim=True):")          #[1, 3, 4]
print(x.mean(axis=0,keepdim=True).shape)
print(x.mean(axis=0,keepdim=True))
#
print("shape of x.mean(axis=0,keepdim=False):")         #[3, 4]
print(x.mean(axis=0,keepdim=False).shape)
print(x.mean(axis=0,keepdim=False))
#
print("shape of x.mean(axis=1,keepdim=True):")          #[2, 1, 4]
print(x.mean(axis=1,keepdim=True).shape)
print(x.mean(axis=1,keepdim=True))
#
print("shape of x.mean(axis=1,keepdim=False):")         #[2, 4]
print(x.mean(axis=1,keepdim=False).shape)
print(x.mean(axis=1,keepdim=False))
#
print("shape of x.mean(axis=2,keepdim=True):")          #[2, 3, 1]
print(x.mean(axis=2,keepdim=True).shape)
print(x.mean(axis=2,keepdim=True))
#
print("shape of x.mean(axis=2,keepdim=False):")         #[2, 3]
print(x.mean(axis=2,keepdim=False).shape)
print(x.mean(axis=2,keepdim=False))

 

shape of x:
torch.Size([2, 3, 4])
shape of x.mean(axis=0,keepdim=True):
torch.Size([1, 3, 4])
tensor([[[ 7.,  8.,  9., 10.],
         [11., 12., 13., 14.],
         [15., 16., 17., 18.]]])
shape of x.mean(axis=0,keepdim=False):
torch.Size([3, 4])
tensor([[ 7.,  8.,  9., 10.],
        [11., 12., 13., 14.],
        [15., 16., 17., 18.]])
shape of x.mean(axis=1,keepdim=True):
torch.Size([2, 1, 4])
tensor([[[ 5.,  6.,  7.,  8.]],

        [[17., 18., 19., 20.]]])
shape of x.mean(axis=1,keepdim=False):
torch.Size([2, 4])
tensor([[ 5.,  6.,  7.,  8.],
        [17., 18., 19., 20.]])
shape of x.mean(axis=2,keepdim=True):
torch.Size([2, 3, 1])
tensor([[[ 2.5000],
         [ 6.5000],
         [10.5000]],

        [[14.5000],
         [18.5000],
         [22.5000]]])
shape of x.mean(axis=2,keepdim=False):
torch.Size([2, 3])
tensor([[ 2.5000,  6.5000, 10.5000],
        [14.5000, 18.5000, 22.5000]])

 

keepdim=True
运算完之后的维度和原来一样，原来是三维数组现在还是三维数组（不过某一维度变成了1）；

keepdim=False
运算完之后一般少一维度，求平均变为1的那一维没有了；

axis=k
按第k维运算，其他维度不遍，第k维变为1

# print(x.mean().shape)
# print(x.mean())

shape of x:
torch.Size([2, 3, 4])
torch.Size([])
tensor(12.5000)#所有值的平均值