一、axis

1.假设我们有一个5行4列的矩阵，见下图所示。
这个矩阵的shape：[5][4],维度是2，axis：0,1，按照行，轴为0，按照列轴为1。
当我们按照axis=0，求sum时：我们得到的shape为[4]
当我们按照axis=1，求sum时：我们得到的shape为[5]

2.技巧：
axis=几，就把对应轴的维度去掉。例如：
假设我们的矩阵shape[2,5,4]
当 axis=1 求和时，得到矩阵的shape[2,4]
当 axis=0 求和时，得到矩阵的shape[5,4]
当 axis=[0,2] 求和时，得到矩阵的shape[5]

二、keepdims

1.keepdims ：在进行某些数组/矩阵操作时的一个参数，通常用于指示操作的结果是否应该保留原始数组的维度信息，默认值为false。
2.假设我们的矩阵shape[2,5,4]
当 axis=1 ,keepdims=True求和时，得到矩阵的shape[2,1,4]
当 axis=2 ,keepdims=True求和时，得到矩阵的shape[2,5,1]
当 axis=[1,2] ,keepdims=True求和时，得到矩阵的shape[2,1,1]

三、代码演示

import torch
a = torch.ones((2,5,4))
a.shape

运行结果：

1.axis

a.sum(axis=1).shape

运行结果：

a.sum(axis=0).shape

运行结果：

a.sum(axis=[0,2]).shape

运行结果：

2.keepdims

a.sum(axis=1,keepdims=True).shape

运行结果：

a.sum(axis=2,keepdims=True).shape

运行结果：

a.sum(axis=[1,2],keepdims=True).shape

运行结果：

☺☺☺☺☺☺☺☺☺☺☺☺☺☺☺☺☺☺☺☺☺☺☺☺☺☺☺☺☺☺☺☺☺☺☺☺☺☺☺☺☺☺☺
我是韩一，用知识认识更多的人，欢迎大家指正！