torch.mean()的使用方法

对一个三维数组的每一维度进行操作

1，dim=0

a = torch.Tensor([0, 1, 2, 3, 4, 5,6,7]).view(2, 2, 2) 
print(a) 
mean = torch.mean(a, 0) 
print(mean, mean.shape)

输出结果：

tensor([[[0., 1.],

             [2., 3.]],

             [[4., 5.],

              [6., 7.]]])

tensor([[2., 3.],

            [4., 5.]]) torch.Size([2, 2])

2，dim=1

a = torch.Tensor([0, 1, 2, 3, 4, 5,6,7]).view(2, 2, 2) 
print(a) 
mean = torch.mean(a, 1) 
print(mean, mean.shape)

输出结果

tensor(

[[[0., 1.],

[2., 3.]],

[[4., 5.],

[6., 7.]]])

tensor(

[[1., 2.],

[5., 6.]]) torch.Size([2, 2])

3，dim=2

a = torch.Tensor([0, 1, 2, 3, 4, 5,6,7]).view(2, 2, 2) 
print(a) 
mean = torch.mean(a, 2) 
print(mean, mean.shape)

输出结果

tensor(

[[[0., 1.],

[2., 3.]],

[[4., 5.],

[6., 7.]]])

tensor(

[[0.5000, 2.5000],

[4.5000, 6.5000]]) torch.Size([2, 2])

补充，如果在函数中添加了True,表示要和原来数的维度一致，不够的用维度1来添加，如下

a = torch.Tensor([0, 1, 2, 3, 4, 5,6,7]).view(2, 2, 2) 
print(a) 
mean = torch.mean(a, 2, True) 
print(mean, mean.shape)
tensor([[[0., 1.],
         [2., 3.]],

        [[4., 5.],
         [6., 7.]]])
tensor([[[0.5000],
         [2.5000]],

        [[4.5000],
         [6.5000]]]) torch.Size([2, 2, 1])

补充多维度变化

a = torch.Tensor([0, 1, 2, 3, 4, 5,6,7,8,9,10,11,12,13,14,15]).view(2, 2, 2,2) 
print(a) 
mean = torch.mean(a, 0, True) 
print(mean, mean.shape)
tensor([[[[ 0.,  1.],
          [ 2.,  3.]],

         [[ 4.,  5.],
          [ 6.,  7.]]],


        [[[ 8.,  9.],
          [10., 11.]],

         [[12., 13.],
          [14., 15.]]]])
tensor([[[[ 4.,  5.],
          [ 6.,  7.]],

         [[ 8.,  9.],
          [10., 11.]]]]) torch.Size([1, 2, 2, 2])

a = torch.Tensor([0, 1, 2, 3, 4, 5,6,7,8,9,10,11,12,13,14,15]).view(2, 2, 2,2) 
print(a) 
mean = torch.mean(a, 1, True) 
print(mean, mean.shape)
tensor([[[[ 0.,  1.],
          [ 2.,  3.]],

         [[ 4.,  5.],
          [ 6.,  7.]]],


        [[[ 8.,  9.],
          [10., 11.]],

         [[12., 13.],
          [14., 15.]]]])
tensor([[[[ 2.,  3.],
          [ 4.,  5.]]],


        [[[10., 11.],
          [12., 13.]]]]) torch.Size([2, 1, 2, 2])

a = torch.Tensor([0, 1, 2, 3, 4, 5,6,7,8,9,10,11,12,13,14,15]).view(2, 2, 2,2) 
print(a) 
mean = torch.mean(a, 2, True) 
print(mean, mean.shape)

tensor([[[[ 0.,  1.],
          [ 2.,  3.]],

         [[ 4.,  5.],
          [ 6.,  7.]]],


        [[[ 8.,  9.],
          [10., 11.]],

         [[12., 13.],
          [14., 15.]]]])
tensor([[[[ 1.,  2.]],

         [[ 5.,  6.]]],


        [[[ 9., 10.]],

         [[13., 14.]]]]) torch.Size([2, 2, 1, 2])

a = torch.Tensor([0, 1, 2, 3, 4, 5,6,7,8,9,10,11,12,13,14,15]).view(2, 2, 2,2) 
print(a) 
mean = torch.mean(a, 3, True) 
print(mean, mean.shape)
tensor([[[[ 0.,  1.],
          [ 2.,  3.]],

         [[ 4.,  5.],
          [ 6.,  7.]]],


        [[[ 8.,  9.],
          [10., 11.]],

         [[12., 13.],
          [14., 15.]]]])
tensor([[[[ 0.5000],
          [ 2.5000]],

         [[ 4.5000],
          [ 6.5000]]],


        [[[ 8.5000],
          [10.5000]],

         [[12.5000],
          [14.5000]]]]) torch.Size([2, 2, 2, 1])

a = torch.Tensor([0, 1, 2, 3, 4, 5,6,7,8,9,10,11,12,13,14,15,0, 1, 2, 3, 4, 5,6,7,8,9,10,11,12,13,14,15]).view(2, 2, 2,2,2) 
print(a) 
mean = torch.mean(a, 3, True) 
print(mean, mean.shape)
tensor([[[[[ 0.,  1.],
           [ 2.,  3.]],

          [[ 4.,  5.],
           [ 6.,  7.]]],


         [[[ 8.,  9.],
           [10., 11.]],

          [[12., 13.],
           [14., 15.]]]],



        [[[[ 0.,  1.],
           [ 2.,  3.]],

          [[ 4.,  5.],
           [ 6.,  7.]]],


         [[[ 8.,  9.],
           [10., 11.]],

          [[12., 13.],
           [14., 15.]]]]])
tensor([[[[[ 1.,  2.]],

          [[ 5.,  6.]]],


         [[[ 9., 10.]],

          [[13., 14.]]]],



        [[[[ 1.,  2.]],

          [[ 5.,  6.]]],


         [[[ 9., 10.]],

          [[13., 14.]]]]]) torch.Size([2, 2, 2, 1, 2])