torch.mean和torch.max函数在二维矩阵上的用法实例

>>> import torch
>>> x = torch.arange(15).view(3,5)*1.0  #乘1.0因为torch.mean只能处理float类型
>>> print(x)
tensor([[ 0.,  1.,  2.,  3.,  4.],
        [ 5.,  6.,  7.,  8.,  9.],
        [10., 11., 12., 13., 14.]])

torch.mean(input, dim, keepdim)

>>> x_mean0 = torch.mean(x, dim=0, keepdim=True)
>>> print(x_mean0)
tensor([[5., 6., 7., 8., 9.]])
>>> 
>>> x_mean1 = torch.mean(x, dim=1 ,keepdim=True)
>>> print(x_mean1)
tensor([[ 2.],
        [ 7.],
        [12.]])
>>> 

torch.max(input, dim, keepdim)

>>> values0, indices0 = torch.max(x, dim=0 ,keepdim=True)
>>> print(values0)
tensor([[10., 11., 12., 13., 14.]])
>>> print(indices0)
tensor([[2, 2, 2, 2, 2]])
>>>
>>> values1, indices1 = torch.max(x, dim=1 ,keepdim=True)
>>> print(values1)
tensor([[ 4.],
        [ 9.],
        [14.]])
>>> print(indices1)
tensor([[4],
        [4],
        [4]])
>>> 

keepdim的作用

>>> x_mean = torch.mean(x, dim=0, keepdim=True)  # keepdim=True保留原维度
>>> print(x_mean)
tensor([[5., 6., 7., 8., 9.]])  # 2维
>>> x_mean = torch.mean(x, dim=0,keepdim=False)  # keepdim=False不保留原维度
>>> print(x_mean)
tensor([5., 6., 7., 8., 9.])  # 1维
>>>