Python 中的向量化运算

这篇文章主要介绍 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. 它们相等
  2. 其中一个是 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 mismatched

4. 向量化的矩阵乘法计算

第四种方法,我们使用 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 循环计算花费时间很长,库函数最快,向量化计算仅次于库函数的速度。

那么这个结论可以给我们哪些启示呢?

  1. 库函数由于有官方的底层优化,效果最好。
  2. 很多情况下,做科研时提出了很多 trick 的计算方式,很难直接使用库函数,这时我们应该将数组完全向量化再做计算,可以节省大量的计算时间。
  3. 完全向量化的计算,能够移植到神经网络框架中,虽然用户接口是 Python 的,但其底层还是 C++ 计算,再加上 GPU 的加速,程序的效率将大大提升。
  4. 程序速度提升了 20 倍,就能比别人多做 20 倍的实验,也就大大增加了科研成果的产生几率。

References

  • https://meijun.github.io/vectorized-computing/
  • https://docs.scipy.org/doc/numpy-1.13.0/user/basics.broadcasting.html
Buy me a coffee
Zhao Li /
Published under (CC) BY-NC-SA in categories Python  tagged with 向量化