numpy sum axis详解

axis

先看懂numpy.argmax的含义.那么numpy.sum就非常好理解.
看一维的例子.

import numpy as np
a = np.array([1, 5, 5, 2])
print(np.sum(a, axis=0))

上面代码就是把各个值加相加.默认axis为0.axis在二维以上数组中才能体现出来作用.

import numpy as np
a = np.array([[1, 5, 5, 2],
              [9, 6, 2, 8],
              [3, 7, 9, 1]])
print(np.sum(a, axis=0))

为了描述方便,a就表示这个二维数组,np.sum(a, axis=0)的含义是a[0][j],a[1][j],a[2]j对应项相加的结果.即[1,5,5,2]+[9,6,2,8]+[3,7,9,1]=[13,18,16,11].接着看axis=1的情况.

import numpy as np
a = np.array([[1, 5, 5, 2],
              [9, 6, 2, 8],
              [3, 7, 9, 1]])
print(np.sum(a, axis=1))

np.sum(a, axis=1)的含义是a[i][0],a[i][1],a[i][2],a[i]3对应项相加的结果.即[1,9,3]+[5,6,7]+[5,2,9]+[2,8,1]=[13,25,20].
三维情况是类似的.

import numpy as np
a = np.array([
              [
                  [1, 5, 5, 2],
                  [9, -6, 2, 8],
                  [-3, 7, -9, 1]
              ],

              [
                  [-1, 5, -5, 2],
                  [9, 6, 2, 8],
                  [3, 7, 9, 1]
              ]
            ])
print(np.sum(a, axis=0))

np.sum(a, axis=0)的含义是a[0][j][k],a[1][j][k] (j=0,1,2,k=0,1,2,3)中对应项相加的结果.[[1, 5, 5, 2],[9, -6, 2, 8],[-3, 7, -9, 1]]+[[-1, 5, -5, 2],[9, 6, 2, 8],[3, 7, 9, 1]]=[[0,10,0,4],[18,0,4,16],[0,14,0,2]]. axis=1,axis=2的道理是类似的.

keepdims

keepdims的含义是是否保持维数,默认是false.通过上面的例子可以发现sum之后3维变成2维.2维变成1维.keepdims=True,最直观的理解就是把sum结果又加一个[],以保持它的维度不变.这在某些场景有非常有用.

import numpy as np
a = np.array([
              [
                  [1, 5, 5, 2],
                  [9, -6, 2, 8],
                  [-3, 7, -9, 1]
              ],

              [
                  [-1, 5, -5, 2],
                  [9, 6, 2, 8],
                  [3, 7, 9, 1]
              ]
            ])
print(np.sum(a, axis=0, keepdims=True))

可以和上面的例子对比下结果.

参考资料

numpy官方文档