深入理解Numpy中sum求和的axis参数

Posted 2022-12-05 算法与编程之美

问题

Numpy是最常见的多维向量处理的工具库，求和则是多维向量的常见操作，然而sum求和方法的axis参数的理解对于多数人而言，特别容易混淆和出错，尤其对于初学者更是困难重重。

关于axis参数的理解，互联网上已经有大量相关介绍其用法的文章，然而这些文章多数描述较为混乱，或是简单的复制官方文档的定义，或是肤浅的理解，对于读者而言，看了比不看更难受，更加难以理解。

针对上述问题，本文提出一种新颖的理解axis参数的方法，可以帮助读者更好的掌握多维向量求和的方法。

方法

首先，需要了解axis参数表示针对哪一个维度进行压缩，例如设有一尺寸为(2,3,4)的多维向量，现要求对axis=0进行压缩，那么将会得到(3,4)尺寸的向量；若对axis=1进行压缩，则得到(2,4)尺寸；若对axis=2进行压缩则得到(2,3)尺寸。由此可见，axis指定了压缩的维度。

其次，需要了解待压缩的维度有几个元素，而元素的个数就是原始尺寸中待压缩的数字，如(2,3,4)的多维向量，需要对axis=1进行压缩，则意味着三个同等尺寸的元素需要合并为一个；

下面分别以二维和三维向量为例讲解axis参数的具体用法。

二维向量示例

首先来看代码如下：

import numpy as np

a = np.array([
    [0, 1],
    [1, 2],
    [2, 3],
])
print(a) # [3,2]

print(np.sum(a, axis=0)) # [3 6]
print(np.sum(a, axis=1)) # [1 3 5]

针对axis=0进行压缩，而0维有3个元素，现需要将3个元素压缩成1个元素，即[3,6];
同理，针对axis=1进行压缩，而1维有2个元素，压缩后就成为[1,3,5]

三维向量示例

import numpy as np


a = np.random.randint(0, 2, size=(2,3,4))
print(a)
print('-'*10)

# 0维有两个(3,4)的Array,将其压缩成一个
print(np.sum(a, axis=0))

# 1维有三个,将其压缩成一个
print(np.sum(a, axis=1))

# 2维有四个元素,将其压缩成一个
print(np.sum(a, axis=2))

结语