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

比如:


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 循环计算花费时间很长,库函数最快,向量化计算仅次于库函数的速度。


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


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

分享给朋友:
您可能感兴趣的文章:
随机阅读: