我一直在尝试优化一个涉及大型多维数组计算的Python代码。但是,我在使用numba时得到了令人费解的结果。我的电脑是2015年中期发布的MBP,2.5 GHz i7 四核处理器,运行OS 10.10.5操作系统和Python 2.7.11版本。请看下面的代码:
import numpy as np
from numba import jit, vectorize, guvectorize
import numexpr as ne
import timeit
def add_two_2ds_naive(A,B,res):
for i in range(A.shape[0]):
for j in range(B.shape[1]):
res[i,j] = A[i,j]+B[i,j]
@jit
def add_two_2ds_jit(A,B,res):
for i in range(A.shape[0]):
for j in range(B.shape[1]):
res[i,j] = A[i,j]+B[i,j]
@guvectorize(['float64[:,:],float64[:,:],float64[:,:]'],
'(n,m),(n,m)->(n,m)',target='cpu')
def add_two_2ds_cpu(A,B,res):
for i in range(A.shape[0]):
for j in range(B.shape[1]):
res[i,j] = A[i,j]+B[i,j]
@guvectorize(['(float64[:,:],float64[:,:],float64[:,:])'],
'(n,m),(n,m)->(n,m)',target='parallel')
def add_two_2ds_parallel(A,B,res):
for i in range(A.shape[0]):
for j in range(B.shape[1]):
res[i,j] = A[i,j]+B[i,j]
def add_two_2ds_numexpr(A,B,res):
res = ne.evaluate('A+B')
if __name__=="__main__":
np.random.seed(69)
A = np.random.rand(10000,100)
B = np.random.rand(10000,100)
res = np.zeros((10000,100))
我现在可以在各种函数上运行timeit:
%timeit add_two_2ds_jit(A,B,res)
1000 loops, best of 3: 1.16 ms per loop
%timeit add_two_2ds_cpu(A,B,res)
1000 loops, best of 3: 1.19 ms per loop
%timeit add_two_2ds_parallel(A,B,res)
100 loops, best of 3: 6.9 ms per loop
%timeit add_two_2ds_numexpr(A,B,res)
1000 loops, best of 3: 1.62 ms per loop
似乎“并行”功能即使使用了单个核心的大部分资源,也无法生效,因为在 top
中显示 Python 对“并行”使用了约 40% 的 CPU,对“cpu”使用了约 100%,而 numexpr 使用了约 300%。
guvectorize
的重点在于,您定义的操作将应用于任何_额外_维度(这是将以并行方式完成的部分)。 您编写的代码本身不会被并行化。 因此,如果您将A
,B
和res
更改为形状为(10000,100,100)
,则来自第三个维度的 100 个不同迭代将并行运行。 - DavidW