torch.mean

torch.mean

官方文档: pytorch

pytorch的torch.mean函数有两种用法: 一种是不带任何参数的，他返回的是tensor中所有元素的均值；第二种是带参数的，他返回某一维度的均值，这里分别介绍

第一种

torch.mean(input)

Returns the mean value of all elements in the input tensor.

返回所有元素的均值

a = torch.rand(2,2)
b = torch.mean(a)

'''
a
tensor([[0.8782, 0.3101],
        [0.9752, 0.9772]])

b
tensor(0.7852)
'''

第二种

torch.mean(input, dim, keepdim=False, *, out=None)

返回对应维度元素的均值

• input: 需要求均值的tensor
• dim: 求均值的维度，参考normalize，第几维度，就是求那一维度的均值，例如tensor维度为(3,4,5)，dim=1，那么就是求第一维度，每一组有四个元素的均值
• keepdim: 是否保持以前的维度数量，而不是大小

下面举两个例子

a = torch.rand((2,3,4))
b = torch.mean(a, dim=1)

'''
a
tensor([[[0.9809, 0.6421, 0.6771, 0.9645],
         [0.5214, 0.8376, 0.8131, 0.7947],
         [0.1607, 0.3469, 0.9644, 0.9500]],

        [[0.1877, 0.3958, 0.1021, 0.7560],
         [0.4673, 0.8111, 0.9186, 0.6773],
         [0.4850, 0.8042, 0.1683, 0.2345]]])

b
tensor([[0.5543, 0.6088, 0.8182, 0.9031],
        [0.3800, 0.6704, 0.3963, 0.5559]])
'''

这个例子可以看出，dim=1计算的是3那一维的最大值，可以看出0.5543=mean(0.9809, 0.5214, 0.1607)，此时维度从(2,3,4)变为了(2,4)
然后我们通过下面这个例子看一下keepdim的作用

c = torch.max(a, dim=1, keepdim=True)

'''
c
tensor([[[0.5543, 0.6088, 0.8182, 0.9031]],
        [[0.3800, 0.6704, 0.3963, 0.5559]]])
'''

c的维度变为了(2,1,4)，发现最后c的维度还是三维，不过求mean的那一维度从3变成了1