pytorch框架中，sum()函数中axis = 的理解

首先创建一个三维张量。

A = torch.arange(20*2,dtype=torch.float32).reshape(2, 5, 4)
A.shape, A.sum(), A

输出的结果为。

(torch.Size([2, 5, 4]),
 tensor(780.),
 tensor([[[ 0.,  1.,  2.,  3.],
          [ 4.,  5.,  6.,  7.],
          [ 8.,  9., 10., 11.],
          [12., 13., 14., 15.],
          [16., 17., 18., 19.]],
 
         [[20., 21., 22., 23.],
          [24., 25., 26., 27.],
          [28., 29., 30., 31.],
          [32., 33., 34., 35.],
          [36., 37., 38., 39.]]]))

可以看到，当sum()函数中没有axis = 时，输出的结果为简单求和。

A_sum_axis0 = A.sum(axis=0)
A_sum_axis0, A_sum_axis0.shape

在sum()函数中添加axis = 0，可以理解为对三维数据中(2 ,5, 4)的2这一维度进行求和

输出结果为。

(tensor([[20., 22., 24., 26.],
         [28., 30., 32., 34.],
         [36., 38., 40., 42.],
         [44., 46., 48., 50.],
         [52., 54., 56., 58.]]),
 torch.Size([5, 4]))

 axis = 1，axis = 2 时同理

(tensor([[ 40.,  45.,  50.,  55.],
         [140., 145., 150., 155.]]),
 torch.Size([2, 4]))

(tensor([[  6.,  22.,  38.,  54.,  70.],
         [ 86., 102., 118., 134., 150.]]),
 torch.Size([2, 5]))