浅谈numpy广播机制

红烧code 2023-03-21 我要评论

Broadcast

广播是numpy对不同形状(shape)的数组进行数值计算的方式，对数组的算术运算通常在相应的元素上进行。如果两个数组a和b形状相同，即a.shape==b.shape为True时，那么a*b的结果就是a与b数组的对应位相乘。这要求维数相同，并且各维度的长度相同。

例如：

import numpy as np

a = np.array([1, 2, 3])
b = np.array([4, 5, 6])
print(a*b)


# 运行结果为array([4,10,18])

最简单的广播机制

当一个常数与一个array对象相乘时，这时就会触发numpy的广播机制

import numpy as np

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


# 运行结果为array([2,4,6])

或许你把这当作理所应当，但是这就是通过广播机制来实现的。

稍微复杂的广播机制

import numpy as np
a = np.array([[0, 0, 0],
              [10, 10, 10],
              [20, 20, 20],
              [30, 30, 30]])
b = np.array([1, 2, 3])
print(a+b)

输出结果是：

array([[ 1, 2, 3],
[11, 12, 13],
[21, 22, 23],
[31, 32, 33]])

初学者可能认为这样的运算会报错，但是这的确是正确的，原因还是广播机制

广播机制到底做了什么

我们观察上面这个稍显复杂的例子：

a.shape=(4,3)而b.shape=(1,3)，两个矩阵axis=1的数据宽度是相同的，但是axis=0的数据宽度不一样，并且b.shape[0]=1，这就是广播机制的触发条件，这时numpy会把b沿axis=0的方向复制4份，即形状变成(4, 3)，与a的一致，接下来就是对应位相加即可。

如下图：

一个典型的错误案例

import numpy as np
a = np.array([[0, 0, 0],
              [10, 10, 10],
              [20, 20, 20],
              [30, 30, 30]])
b = np.array([0, 1, 2, 3])
a+b

提示无法进行广播：

能否广播必须从axis的最大值向最小值看去，依次对比两个要进行运算的数组的axis的数据宽度是否相等，如果在某一个axis下，一个数据宽度为1，另一个数据宽度不为1，那么numpy就可以进行广播；但是一旦出现了在某个axis下两个数据宽度不相等，并且两者全不为1的状况，就无法广播。

一个正确的经典示例

import numpy as np
a = np.arange(6).reshape((2, 3, 1))
b = np.arange(6).reshape((1, 3, 2))
print(a)
print(b)
print(a+b)
print((a+b).shape)

结果：

a是：
array([[[0],
[1],
[2]],

[[3],
[4],
[5]]])

b是：
array([[[0, 1],
[2, 3],
[4, 5]]])

a+b是：
array([[[ 0, 1],
[ 3, 4],
[ 6, 7]],

[[ 3, 4],
[ 6, 7],
[ 9, 10]]])

运算过程是：

1不等于2，所以对a的axis=2进行广播，a变为

array([[[0,0],
        [1,1],
        [2,2]],

       [[3,3],
        [4,4],
        [5,5]]])

axis=1时3等于3，不进行广播

axis=0时，2不等于1，所以对b的axis=0进行广播，b变为：

array([[[0, 1],
        [2, 3],
        [4, 5]],
  
       [[0, 1],
        [2, 3],
        [4, 5]]])

最后a、b的shape变为一样，运算结果即为：

array([[[ 0,  1],
        [ 3,  4],
        [ 6,  7]],

       [[ 3,  4],
        [ 6,  7],
        [ 9, 10]]])

一种更便捷的计算方式

注：np.newaxis表示增加一个数据宽度为1的维度

因此a[:, np.newaxis] = array([[0],[10],[20],[30]])

import numpy as np
a = np.array([0, 10, 20, 30])
b = np.array([0, 1, 2])
a[:, np.newaxis]+b

结果：

array([[ 0, 1, 2],
[10, 11, 12],
[20, 21, 22],
[30, 31, 32]])

广播过程：