这篇文章主要介绍 Python 中数值计算的不同方法与性能对比，以矩阵乘法为例，探索加快模型训练的方法。
1. 使用 Python for 循环计算矩阵乘法
第一种方法是直接使用 Python 的循环来计算矩阵乘法，matmul函数中计算矩阵 B(m*r) 和矩阵 C(r*n) 相乘的结果。
我们先让程序运行一次乘法来 warm up，避免指令的初始运行带来的时间误差，然后统计矩阵纬度在 0~500 之间变化时，矩阵相乘所花费的时间（下同）。
import sys
import time
import random
def matmul(a, b, c, m, n, r):
for i in range(m):
for j in range(n):
for k in range(r):
a[i][j] += b[i][k]*c[k][j]
n = int(sys.argv[1])
a = [[0.0 for _ in range(n)] for _ in range(n)]
b = [[random.random() for _ in range(n)] for _ in range(n)]
c = [[random.random() for _ in range(n)] for _ in range(n)]
matmul(a, b, c, n, n, n) # warm up
a = [[0.0 for _ in range(n)] for _ in range(n)]
b = [[random.random() for _ in range(n)] for _ in range(n)]
c = [[random.random() for _ in range(n)] for _ in range(n)]
tic = time.time()
matmul(a, b, c, n, n, n)
toc = time.time()
print(toc-tic)N = 300 时的花费时间为 5.774713993s。
2. 使用 NumPy for 循环计算矩阵乘法
第二种方法使用的是 NumPy 的数据结构。
def matmul(a, b, c, m, n, r):
for i in range(m):
for j in range(n):
for k in range(r):
a[i][j] += b[i][k]*c[k][j]
a = np.zeros((n, n))
b = np.random.rand(n, n)
c = np.random.rand(n, n)N = 300 时的花费时间为 22.72896314s。
3. 使用库函数的矩阵乘法
第三种方法使用 NumPy 自带的矩阵乘法操作。
def matmul(a, b, c, m, n, r):
a = np.matmul(b, c)N = 300 时的花费时间为 0.002428055s。
在介绍最后一种方法时，我们先来学习一下 NumPy 的广播机制。
NumPy 的 Broadcasting 机制
观察这个例子：
>>> B = np.array([[0, 1]]) # shape: 1x2
>>> C = np.array([[0],
... [1],
... [2]]) # shape: 3x1
>>> B + C
array([[0, 1],
[1, 2],
[2, 3]])Broadcasting 是 NumPy 中对形状不同的两个数组做运算的机制，较小的数组会沿着较大的数组“广播”，这样它们的纬度能够兼容。
值得一提的是，Broadcasting 是 element-wise 的，也就是一对一做运算的，而不是像矩阵乘法那样。
两个数组形状一样：
>>> a = np.array([1.0, 2.0, 3.0])
>>> b = np.array([2.0, 2.0, 2.0])
>>> a * b
array([ 2., 4., 6.])数组与单个值的计算：
>>> a = np.array([1.0, 2.0, 3.0])
>>> b = 2.0
>>> a * b
array([ 2., 4., 6.])两种计算方式中，广播要更高效，因为它在乘法计算时动用更少的内存。
Broadcasting 规则
对两个数组进行操作时，NumPy 从后到前比较形状。两个维度兼容的条件是：
- 它们相等
- 其中一个是 1
比如：
A (2d array): 5 x 4
B (1d array): 1
Result (2d array): 5 x 4
A (2d array): 5 x 4
B (1d array): 4
Result (2d array): 5 x 4
A (3d array): 15 x 3 x 5
B (3d array): 15 x 1 x 5
Result (3d array): 15 x 3 x 5
A (3d array): 15 x 3 x 5
B (2d array): 3 x 5
Result (3d array): 15 x 3 x 5
A (3d array): 15 x 3 x 5
B (2d array): 3 x 1
Result (3d array): 15 x 3 x 5不满足规则的情况：
A (1d array): 3
B (1d array): 4 # trailing dimensions do not match
A (2d array): 2 x 1
B (3d array): 8 x 4 x 3 # second from last dimensions mismatched4. 向量化的矩阵乘法计算
第四种方法，我们使用 NumPy 将矩阵完全向量化，然后利用广播机制计算出矩阵乘法的结果，避免了使用循环。
def matmul(a, b, c, m, n, r):
m_r_n = b.reshape(m, r, 1) * c.reshape(1, r, n)
a = np.sum(m_r_n, axis=1)N = 300 时的花费时间为 0.14920187s。
结论
四种方法后，我们做了一个统计：矩阵纬度从 0 到 500 增加时，所花费的计算时间如图：
根据这个统计可以发现，for 循环计算花费时间很长，库函数最快，向量化计算仅次于库函数的速度。
那么这个结论可以给我们哪些启示呢？
- 库函数由于有官方的底层优化，效果最好。
- 很多情况下，做科研时提出了很多 trick 的计算方式，很难直接使用库函数，这时我们应该将数组完全向量化再做计算，可以节省大量的计算时间。
- 完全向量化的计算，能够移植到神经网络框架中，虽然用户接口是 Python 的，但其底层还是 C++ 计算，再加上 GPU 的加速，程序的效率将大大提升。
- 程序速度提升了 20 倍，就能比别人多做 20 倍的实验，也就大大增加了科研成果的产生几率。