让python快到飞起-numba加速

以下文章来源于气海同途 ，作者气海同途

一、前言

python是一门高效动态编程语言，由于其采用简洁明了的语法以及灵活性深受大家欢迎。但是，这既是它最大的优势，也是最大的劣势。它的灵活性和无类型的高级语法可能会导致数据和计算密集型程序的性能不佳，因为运行本地编译代码要比运行动态解释代码快很多倍。

因此，注重效率的 Python 程序员通常会使用 C 语言重写最内层的循环，然后从 Python 调用已编译的 C 语言函数。许多项目都力求简化这种优化（例如 Cython），但它们通常需要学习新的语法。虽然 Cython 显著提高了性能，但可能需要对 Python 代码进行艰巨的手动修改工作。

对于不了解C、C++、Cython等高效语言，而重新学习一门语言的成本又太高的用户而言，Numba 被视作为最佳的替代方案，学习应用要简单得多。无需学习新的语法，也无需替换 Python 解释器、运行单独的编译步骤或安装 C/C++ 编译器。只需将 @jit Numba 修饰器应用于 Python 函数即可。

Numba执行图

Numba 能够动态编译代码，这意味着还可以享受 Python 带来的灵活性。此外，Python 程序中由 Numba 编译的数值算法，可以接近使用编译后的 C 语言或 FORTRAN 语言编写的程序的速度；并且与原生 Python 解释器执行的相同程序相比，运行速度最多快 100 倍。

二、numba的安装：

conda install numba

或者：

pip install numba

三、numba的使用：

我们只需要在原来的代码上添加一行@jit(nopython=True)，即可将一个函数编译成机器码，其他地方都不需要更改。下面以一个简单的案例，做循环计算，来测试numba的加速情况:

from numba import jit
import numpy as np
import time

def cal():    
    x=0    
    for i in np.arange(100000000):        
        x+=i    
    return x
start_time=time.time()
cal()
end_time=time.time()
print('numpy用时：',end_time-start_time,'秒')

#将需要加速的部分封装成函数，在函数前加上numba即时编译装饰器
@jit(nopython=True)
def cal_numba():    
    x=0    
    for i in np.arange(100000000):        
        x+=i    
    return x
start_time=time.time()
cal_numba()
end_time=time.time()
print('numba用时：',end_time-start_time,'秒')

通过测试，未加速的代码用时：16.72s；numba加速后用时：0.6334s，加速效果比较明显。细心的读者可能发现，这里测试使用了1亿次的迭代计算，其实在海洋中这样的计算量并不算大，相当于1000*1000的矩阵100次计算量。下图列出来不同计算量情况下，加速前后的用时：

从测试来看，当计算量大于100万时，numba加速较为明显。对于三维气象海洋数据，100*100*100（时间,纬度,经度）计算1次以上，numba即可有明显的加速效果。

注意两点：

• 使用Numba时，总时间 = 编译时间 + 运行时间。相比所能节省的计算时间，编译的时间开销很小，才能达到加速效果。对于一个需要多次调用的Numba函数，只需要编译一次，后面再调用时就不需要编译了。
• 这里装饰的函数调用的API是有限制的！比如pandas是更高层次的封装，Numba其实不能理解它里面做了什么，所以无法对其加速。一些大家经常用的机器学习框架，如scikit-learn，tensorflow，pytorch等，已经做了大量的优化，不适合再使用Numba做加速。