Numpy大规模矩阵运算优化加速技巧

技巧一:将大规模向量、矩阵运算交给numpy

如果对数组进行向量化运算,例如全体四则运算矩阵乘法求和按指标求和等,一定要利用numpy的矩阵乘法dot和einsum。

  • dot 二维矩阵乘法
    numpy的矩阵运算的王牌,做矩阵乘法的首选,优化到了极致。

  • einsum 一般矩阵乘法
    仅次于dot,比numpy的sum、inner、outer、kron都要快一个或者几个数量级。

夸张的是,einsum求和比向量四则运算都要快:

1
2
3
4
5
a = np.random.random(5000000)
b = np.random.random(5000000)
print(np.isclose(np.einsum('i->',a-b), np.einsum('i->',a) - np.einsum('i->',b)))
%timeit np.einsum('i->',a-b)
%timeit np.einsum('i->',a) - np.einsum('i->',b)

输出:

1
2
3
True
17.8 ms ± 284 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
5.34 ms ± 72.8 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

当数据量大时,两次einsum一次标量减法甚至快于一次einsum一次向量减法。

关于如何将各种专门的矩阵运算转化为einsum及其性能比较,参考链接Why is numpy’s einsum faster than numpy’s built in functions?

技巧二:避免大量重复调用numpy函数解决小规模问题

涉及到向量运算的计算过程,若能够使用numpy矩阵运算函数一次计算,则千万不要再使用for循环。但是有的情况不可避免使用for循环且循环次数巨大,这种情况应当尽量避免在for中调用numpy函数进行小规模运算,例如每一次循环中应该避免

  • 使用all、any或者isclose比较两个二、三元向量的分量
    直接各自比较即可。
  • 使用random下的函数产生单个随机数
    最好在循环外产生一个随机数narray。
  • 每次建立一样的数组(如用来掩模的数组)
    循环外建立,循环内调用即可。
  • 小规模数组合并:array((x,y))
    尽量省去,除非有运算上的必要。
  • 改变narray数据个数
    改变narray长度需要重新申请内存并复制,非常慢,哪怕只加入一个数据。作为替代,频繁添加数据可以用Python的list.append,或者提前申请好narray内存每次直接赋值即可。由于list是动态链表,改变表长是很容易的(至于用栈,就算是深度优化了,在Python中深度优化意义不大)。

技巧三:优先对较低指标求和

1
2
3
4
5
6
a = np.random.rand(500, 500, 500)
%timeit a[0,:,:].sum()

%timeit a[:,0,:].sum()

%timeit a[:,:,0].sum()

输出:

1
2
3
59.5 µs ± 478 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
119 µs ± 2.32 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
4.88 ms ± 230 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

耗时相差几个数量级。这是由于narray数据结构导致的。所以对于需要大量求和运算的指标尽量放在第一位。

技巧四:使用ravel对矩阵一维化

ravel将多维数组一维化时避免拷贝,所以速度较快。相比之下,flatten一维化返回拷贝所以慢几个数量级,reshape(-1)一维化也可以避免拷贝,但是相比ravel要慢一些。

技巧五:使用矩阵自运算替代运算后赋值

自加、自减等运算要比运算后赋值要快,即a*=2要快于a=a*2。因为不需要重新申请内存。

技巧六:一维数组合并

numpy没有append,因为array的数据结构是块状,改变元素个数需要重新申请内存并复制,所以应当尽量减小数组合并的次数。对于必要的合并,concatenate是比较快速的函数,其次可以用r_[],大约慢一个数量级。转化成list再合并是非常慢的。

技巧七:避免花式索引

花式索引(fancy indexing)即用narray作为index取出一个小数组。因为会效率低而且占内存,不适合大量调用。如果需要可以用布尔掩模或者使用take、compress替代。

技巧八:避免narray.T转置数组

转置涉及到拷贝,非常耗时。一个例子是:在向量或者矩阵乘法中需要用到转置,这种场合可以事先将矩阵定义为良好的形状,或者使用einsum替代dot,付出较小的代价。

容易犯错之处:你的原始narray改变了吗?

将一个narray变量a(或者其切片)赋给另一个变量b之后,对b进行了操作,那么a改变了吗?有的时候我们希望b只是a的一个别名,通过修改b来修改a;但有的时候我们希望原始的a不受干扰。问题是:什么时候原始a会变?什么时候不变?

一、修改元素还是修改整体引用

假如我们希望将一个narray变量a的一部分a[0]赋给另一个变量b,并希望通过修改b来改变a。

原始数组什么时候变:

当我们修改b的元素时,会改变a。

1
2
3
4
a = np.zeros((2,2))
b = a[0]
b[0] = 1
print(a)

输出:

1
2
[[ 1.  0.]
[ 0. 0.]]

注意:当b数组的元素不是“数值”而是一个narray的时候,修改b元素即改变b中该narray的指向,同样会改变a。

1
2
3
4
a = np.zeros((2,2))
b = a
b[0] = np.ones(2)
print(a)

输出:

1
2
[[ 1.  1.]
[ 0. 0.]]

原始数组什么时候不变:

假如我们希望将b的所有元素都改成1,那么直接将另一个narray赋给b将不能改变a。

1
2
3
4
a = np.zeros((2,2))
b = a[0]
b = np.ones(10)
print(a)

输出:

1
2
[[ 0.  0.]
[ 0. 0.]]

希望达到修改a的目的的话,则可以通过对b全部赋值实现。

1
2
3
4
a = np.zeros((2,2))
b = a[0]
b[:] = np.ones(2)
print(a)

输出:

1
2
[[ 1.  1.]
[ 0. 0.]]

二、自运算还是运算后赋值

表面上自运算b+=1和运算后赋值b=b+1对b效果是一样的,但是实际上对原始数组a的影响是不一样的。

原始数组什么时候变:

自运算会改变原始数组,因为是原位运算,修改了元素的值。

1
2
3
4
a = np.zeros((2,2))
b = a[0]
b += 1
print(a)

输出:

1
2
[[ 1.  1.]
[ 0. 0.]]

原始数组什么时候不变:

运算后赋值不会改变原始数组,因为是运算后将新数组的引用赋给了b。

1
2
3
4
a = np.zeros((2,2))
b = a[0]
b = b + 1
print(a)

输出:

1
2
[[ 0.  0.]
[ 0. 0.]]

当然全部赋值是允许的,可以强行改变a。

1
2
3
4
a = np.zeros((2,2))
b = a[0]
b[:] = b + 1
print(a)

输出:

1
2
[[ 1.  1.]
[ 0. 0.]]

三、函数传值:python函数传递array时以指针传递

函数传值对被传入narray(实参)的影响,本质上与上面相同。

原始数组什么时候变:

函数传值,相当于将实参赋给形参b=a。假如在函数内改变形参narray元素的值(不仅包括修改“数值”,也包括修改子narray的指向),则实参的元素的值也被改变。

改变数值
1
2
3
4
5
6
def func(a):
a[0] = 1
s = np.zeros(3)
print(s)
func(s)
print(s)

输出:

1
2
[ 0.  0.  0.]
[ 1. 0. 0.]

改变子narray
1
2
3
4
5
6
def func(a):
a[0]=np.ones(10)
b=np.zeros((10,10))
print(b)
func(b)
print(b)

输出:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
[[ 0.  0.  0.  0.  0.  0.  0.  0.  0.  0.]
[ 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
[ 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
[ 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
[ 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
[ 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
[ 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
[ 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
[ 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
[ 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]]
[[ 1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
[ 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
[ 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
[ 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
[ 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
[ 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
[ 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
[ 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
[ 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
[ 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]]

原始数组什么时候不变

在函数内,假如没有修改形参的内容,而是直接将整个形参变量赋为另一个变量,则实参不改变。

直接改变形参
1
2
3
4
5
6
7
def func(a):
a=np.ones(10)
print("in func a=",a)
b=np.zeros(10)
print(b)
func(b)
print(b)

输出:

1
2
3
[ 0.  0.  0.  0.  0.  0.  0.  0.  0.  0.]
in func a= [ 1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
[ 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]

四、在block内新建的narray在block外同样有效

1
2
3
4
c = 1
if c>0:
data = np.ones(3)
print(data)

输出

1
[ 1.  1.  1.]

意味着不需要在if、for、while等block外事先声明或者划分内存。注意假如上面的c=0,则会因data未定义报错。这一点与C++变量的作用域截然不同。

技巧X 使用cProfile测试分析函数运行耗时

1
$ python -m cProfile -s cumulative filename.py > log

从生成的log文件就可以分析各个函数的调用次数和运行时间。

优先优化占用时间最长、调用次数最高的函数,若某个函数的累积运行时间没有达到总运行时间的10%,就没有对其优化的必要。

参考资料

  1. 伯乐在线, 如何获得NumPy的最佳性能