Numba 简介

Numba 是 Python 的一个 JIT (just-in-time) 编译器，最适用于 NumPy 数组、函数，以及 Python 循环。基本上，用法就是给原来的 Python 函数加一个修饰器，当运行到经 Numba 修饰的函数时，它会被编译为机器码，之后再调用时，就能以机器码的速度来执行了。

按我上手使用的经验来看，Numba 对原代码的改动不是太大，对能加速的部分，加速效果明显；对不支持的加速的 Python 语句/第三方库，可以选择不使用 numba 来规避。这是我选择 Numba 的原因。

首先：应该编译(优化)什么？

由于 Numba 本身的限制(稍后介绍)，不能做到对整个程序完全的优化。实际上，也没必要这样做——只需要优化真正耗时间的部分即可。

怎么找到真正耗时间的部分？除了靠直觉，还可以借用工具来分析，例如 Python 自带的 cProfile，还有 line_profiler 等，这里不再细讲。

安装

可以通过 conda 或 pip，一个命令安装：

conda / pip install numba

什么样的代码能加速？

按照官方文档的示例代码，如果代码中含有很多数学运算、使用 NumPy，或者有不少 Python 的 for 循环(Python 性能大忌)，那么 Numba 就能给你很好的效果。尤其是多重 for 循环，可以获得极大的加速。

例如，下面这段代码，就能享受到 JIT：

from numba import jit

import numpy as np

x = np.arange(100).reshape(10, 10)

@jit(nopython=True) # 设置为"nopython"模式 有更好的性能

def go_fast(a): # 第一次调用时会编译

trace = 0

for i in range(a.shape[0]): # Numba likes loops

trace += np.tanh(a[i, i]) # Numba likes NumPy functions

return a + trace # Numba likes NumPy broadcasting

print(go_fast(x))

但是，类似下面的代码，Numba 就没什么效果：

from numba import jit

import pandas as pd

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

@jit

def use_pandas(a): # 这个函数就加速不了

df = pd.DataFrame.from_dict(a) # Numba 不支持 pd.DataFrame

df += 1 # Numba 也不支持这个

return df.cov() # 和这个

print(use_pandas(x))

总之，Numba 应付不了 pandas。以我的经验，需要把 DataFrame 转成 np.ndarray，再输入给 Numba。

要强制用 nopython 模式

刚才 work 的代码中，@jit(nopython=True) 这里传入了 nopython 这个参数，而不 work 的代码中，就没有这个参数。为什么呢？

这是因为，@jit 实际上有两种模式，分为别 nopython 和 object 模式。只有 nopython 模式，才是能真正大幅加速的模式。而 nopython 模式只支持部分的 Python 和 NumPy 函数，如果运行时用到了不支持的函数/方法，程序就会崩掉 (例如刚才不能加速的例子) 。如果不强制设定 nopython 模式，编译函数失败时，会回退到 object 模式，程序虽然不会崩，但却偏离了我们给它加速的本意。

我既然用了 Numba，我就希望它能真正地发挥作用。所以我选择强制开启 nopython ，如果不能加速，不如让它直接抛 exception。

不支持哪些 features?

说了那么多，哪些能用，哪些不能用？当前版本(本文撰写时为 0.45)的 Numba 不支持以下 Python 功能：

Exception handling (try .. except, try .. finally)

Context management (the with statement)

Dict/set/generator comprehensions

Generator delegation (yield from)

此外，还有一些限制，例如 list/tuple 的元素必须同类(静态)等。对于 NumPy，其支持的 feature 见文末的参考链接。

一些使用心得

再次提醒，不支持 pandas，不然分分钟报错。另外发现 pandas 的效率其实挺低的，例如取 column 的操作。

变量不能做类型转换(类似静态语言的要求)。

待加速的函数，要么是运行很多很多次，要么是运行时间很长；这是因为第一次编译都需要时间，如果编译完只运行这个函数几次，说不定还没有直接跑快；虽然有一个可以把编译结果存在硬盘的cache参数，但是实测发现从硬盘读取这个缓存结果也还是需要一点时间的。

涉及到具体的业务代码，如果有 class，或者比较长的函数，最好把真正需要优化的部分提取出来做 JIT，绕过不支持的 methods。class 虽然有初步的支持 (jitclass)，但是限制很多，一般只能用于简单的 class。

相关参考

本文只是很初步的介绍，更深度的使用方法，还是要自己看文档：

Views:

599