python: numba提升程序效率

简介

是什么
numba是一个的即时编译器(just-in-time compiler)。即它能在程序运行时,使用LLVM编译器库将Python函数转换为优化的机器码。从而使得运行速度可以达到C或FORTRAN的速度。

怎么用
只需用numba去修饰python函数即可。换言之，不需要替换Python解释器、运行单独的编译步骤、安装C/ c++编译器等额外操作。

实例1

问题描述
对于矩阵连乘:$U_n...U_2{U}_1\rho$, 原方法是用for循环累乘.测试下用numba.jit修饰后效率是否有提升.

使用注意: array的数据类型必须一样,否则会报错.

实例代码

def f(U, rho):
    for i in range(len(U)):
        rho = np.dot(U[i] ,rho)
        
@numba.jit(nopython=True)        
def f_jit(U, rho):
    for i in range(len(U)):
        rho = np.dot(U[i] ,rho)

a = np.array([[0.1,0.2],[0.3,0.4]])
U = np.array([a]*500)
rho = np.array([[0.5,0.5],[0.5,0.5]])
%timeit f(gates, rho)
%timeit f_jit(gates, rho)

结果分析
对这个测试实例,提升近3倍.

114 µs ± 1.31 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
38.6 µs ± 1.09 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

实例2

问题描述: 矩阵累加,测试用numba后,是否有效率提升.

实例代码

def f(A):
    C = 0
    for a in A:
        C = C + a
    return C

@numba.jit(nopython=True)
def f_jit(A):
    C = np.zeros(A[0].shape,dtype=A.dtype)
    for a in A:
        C = C + a
    return C

A=[]
for i in range(1000):
    a = np.random.uniform(-1,1,(16,16))
    A.append(a)
A = np.asarray(A)
%timeit f(A)
%timeit f_jit(A)

结果分析
对这个测试实例,提升1-2倍

2.81 ms ± 44.6 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
1.05 ms ± 38 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)