Numba加速Python程序

分手后的思念是犯贱 2022-10-16 07:52 154阅读 0赞

Python和Java一样是基于虚拟机的语言，并不是像C/C++那样将程序代码编译成机器语言再运行，而是解释一行执行一行，速度比较慢。使用Numba库的JIT技术编译以后，可以明显提高程序的运行速度。

import numpy as np
    from numba import jit
    import time
    
    @jit
    def sum_jit(arr):
        s_time = time.time()
        m = arr.shape[0]
        result = 0.0
        for i in range(m):
            result += arr[i]
        e_time = time.time()
        return (e_time-s_time)
    
    def sum(arr):
        s_time = time.time()
        m = arr.shape[0]
        result = 0.0
        for i in range(m):
            result += arr[i]
        e_time = time.time()
        return (e_time-s_time)
    
    def main():
        n = int(10.0*1e6)
        array = np.random.random(n)
        t1 = sum_jit(array)
        t2 = sum(array)
        print("Time with JIT:", t1)
        print("Time without JIT:", t2)
    
    if __name__ == '__main__':
        main()

上面代码的sum()和sum\_jit()两个函数完全相同，区别在于sum\_jit()添加了@jit注解，表面该函数使用Numba JIT进行编译以后再运行，而sum()是不编译的运行。对main()中n的值取1×106~10×106，计算每个值两个函数运行的时间，单位为Second，如下图所示：

![watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L21hcnRpbmtlaXRo_size_16_color_FFFFFF_t_70][]

可以看到随着运算量的线性增加，sum()的运行时间几乎也是线性增加的，而sum\_jit()的运行时间基本上保持不变，运算量的增加并没有带来影响。

### 更多numba的加速选项  [https://zhuanlan.zhihu.com/p/60994299][https_zhuanlan.zhihu.com_p_60994299] ###

除了上面提到的jit和vectorize，其实numba还支持很多加速类型。常见的比如

@nb.jit(nopython=True,fastmath=True) 牺牲一丢丢数学精度来提高速度

@nb.jit(nopython=True,parallel=True) 自动进行并行计算

切记一定要用nopython。默认都是True的，但有时候如果定义的函数中遇到numba支持不良好的部分，它就会自动关闭nopython模式。没有nopython的numba就好像没有武器的士兵，虽然好过没兵，但确实没什么战斗力。因此，**在使用jit时候要明确写出nopython=True**。如果遇到问题，就找到这些支持不良好的部分，然后改写。毕竟numba对loop非常友好，改写这些部分应当是非常容易的。

[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L21hcnRpbmtlaXRo_size_16_color_FFFFFF_t_70]: /images/20221014/b1f2d43563dc43a8bf294bdd021a8979.png
[https_zhuanlan.zhihu.com_p_60994299]: https://zhuanlan.zhihu.com/p/60994299