强化学习技巧五：numba提速python程序

Posted 2021-08-08 汀、

numba是一款可以将python函数编译为机器代码的JIT编译器，经过numba编译的python代码（仅限数组运算），其运行速度可以接近C或FORTRAN语言。

numba使用情况

• 使用numpy数组做大量科学计算时
• 使用for循环时

1.numba使用

导入numpy、numba及其编译器

import numpy as np
import numba
from numba import jit

传入numba装饰器jit，编写函数

# 使用numba的情况
@jit(nopython=True)
def t():
    x = 0
    for i in np.arange(5000):
        x += i
    return x
%timeit(t())

nopython = True选项要求完全编译该函数（以便完全删除Python解释器调用），否则会引发异常。这些异常通常表示函数中需要修改的位置，以实现优于Python的性能。强烈建议您始终使用nopython = True。

2.不适用场景

numba目前只支持Python原生函数和部分Numpy函数，其他场景下无效。

from numba import jit
import pandas as pd

x = {'a': [1, 2, 3], 'b': [20, 30, 40]}

@jit
def use_pandas(a): # Function will not benefit from Numba jit
    df = pd.DataFrame.from_dict(a) # Numba doesn't know about pd.DataFrame
    df += 1                        # Numba doesn't understand what this is
    return df.cov()                # or this!

print(use_pandas(x))

上述代码中使用了Pandas，而Pandas并不是原生代码，而是更高层次的封装，Numba不能理解pandas内部在做什么，所以无法对其加速。

而一些常用的机器学习框架，比如scikit-learn, tensorflow, pyrorch等，已经做了大量的优化，不适合再使用Numba做加速。

可以简单总结为，Numba不支持：

• pandas

• scikit-learn, tensorflow, pyrorch

• try…except 异常处理

• with 语句

• yield from

Numba有两种模式：

@jit：object模式：上图左侧
Numba的@jit装饰器会尝试优化代码，如果发现不支持（比如pandas等），那么Numba会继续使用Python原来的方法去执行该函数。

@jit(nopython=True)或者@njit：nopython模式：上图右侧
强制加速，不会进入上图左侧流程，只进行右侧流程，如果编译不成功，就抛出异常。
 

 Numba使用了LLVM和NVVM技术，此技术将Python等解释型语言直接翻译成CPU、GPU可执行的机器码。

那如何决定是否使用Numba呢？

Numba的@jit装饰器就像自动驾驶，用户不需要关注到底是如何优化的，Numba去尝试进行优化，如果发现不支持，那么Numba会继续用Python原来的方法去执行该函数，即图 Python解释器工作原理中左侧部分。这种模式被称为object模式。前文提到的pandas的例子，Numba发现无法理解里面的内容，于是自动进入了object模式。object模式还是和原生的Python一样慢，还有可能比原来更慢。

Numba真正牛逼之处在于其nopython模式。将装饰器改为@jit(nopython=True)或者@njit，Numba会假设你已经对所加速的函数非常了解，强制使用加速的方式，不会进入object模式，如编译不成功，则直接抛出异常。nopython的名字会有点歧义，我们可以理解为不使用很慢的Python，强制进入图 Python解释器工作原理中右侧部分。

实践上，一般推荐将代码中计算密集的部分作为单独的函数提出来，并使用nopython方式优化，这样可以保证我们能使用到Numba的加速功能。其余部分还是使用Python原生代码，在计算加速的前提下，避免过长的编译时间。（有关编译时间的问题下节将会介绍。）Numba可以与NumPy紧密结合，两者一起，常常能够得到近乎C语言的速度。尽管Numba不能直接优化pandas，但是我们可以将pandas中处理数据的for循环作为单独的函数提出来，再使用Numba加速。