Numba 的 ~5 分钟指南

Numba 是一个针对 Python 的即时编译器,它对使用 NumPy 数组和函数以及循环的代码效果最佳。使用 Numba 最常见的方式是通过其一系列装饰器,这些装饰器可以应用于你的函数,以指示 Numba 编译它们。当调用一个由 Numba 装饰的函数时,它会被编译为机器代码以“即时”执行,并且你的代码的全部或部分可以随后以原生机器代码的速度运行!

开箱即用,Numba 支持以下内容:

  • 操作系统:Windows(64位),OSX,Linux(64位)。*BSD 系统上的非官方支持。

  • 架构:x86, x86_64, ppc64le, armv8l (aarch64), M1/Arm64。

  • GPUs: Nvidia CUDA.

  • CPython

  • NumPy 1.22 - 1.26

我该如何获取它?

Numba 作为 conda 包,适用于 Anaconda Python 发行版:

$ conda install numba

Numba 也有可用的轮子:

$ pip install numba

Numba 也可以 从源代码编译 ,尽管我们不建议初次使用 Numba 的用户这样做。

Numba 通常作为一个核心包使用,因此其依赖性保持在绝对最小,然而,可以通过以下方式安装额外的包以提供附加功能:

  • scipy - 支持编译 numpy.linalg 函数。

  • colorama - 支持在回溯/错误信息中进行颜色高亮显示。

  • pyyaml - 通过 YAML 配置文件启用 Numba 的配置。

  • intel-cmplr-lib-rt - 允许使用 Intel SVML(高性能短向量数学库,仅限 x86_64)。安装说明在 性能提示 中。

Numba 适用于我的代码吗?

这取决于你的代码是什么样的,如果你的代码是数值导向的(做了很多数学运算),大量使用 NumPy 和/或有很多循环,那么 Numba 通常是一个不错的选择。在这些例子中,我们将应用 Numba 最基本的 JIT 装饰器 @jit,试图加速一些函数,以展示什么效果好,什么效果不好。

Numba 在类似这样的代码上表现良好:

from numba import jit
import numpy as np

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

@jit
def go_fast(a): # Function is compiled to machine code when called the first time
    trace = 0.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))

如果代码看起来像这样,它将无法很好地工作,甚至根本无法工作:

from numba import jit
import pandas as pd

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

@jit(forceobj=True, looplift=True) # Need to use object mode, try and compile loops!
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 不被 Numba 理解,因此 Numba 将通过解释器运行此代码,但会增加 Numba 内部开销的成本!

什么是 对象模式

Numba 的 @jit 装饰器基本上在两种编译模式下运行,nopython 模式和 object 模式。在上面的 go_fast 示例中,@jit 装饰器默认在 nopython 模式下运行。nopython 编译模式的行为是基本上编译装饰函数,使其完全不依赖于 Python 解释器运行。这是使用 Numba jit 装饰器的推荐和最佳实践方式,因为它能带来最佳性能。

如果 nopython 模式下的编译失败,Numba 可以使用 对象模式 进行编译。这是通过在 @jit 装饰器中使用 forceobj=True 关键字参数实现的(如上面的 use_pandas 示例所示)。在这种模式下,Numba 会假设所有内容都是 Python 对象,并在解释器中运行代码。指定 looplift=True 可能会比纯 对象模式 获得一些性能提升,因为 Numba 会尝试将循环编译成在机器代码中运行的函数,并在解释器中运行其余代码。为了获得最佳性能,通常应避免使用 对象模式

如何测量 Numba 的性能?

首先,回想一下,Numba 必须在执行函数的机器代码版本之前为给定的参数类型编译您的函数。这需要时间。然而,一旦编译完成,Numba 会为特定类型的参数缓存函数的机器代码版本。如果再次使用相同类型的参数调用它,它可以重用缓存的版本,而不必再次编译。

在测量性能时,一个非常常见的错误是没有考虑到上述行为,并且只用一个简单的计时器对代码计时一次,这个计时器包括了编译函数所花费的时间在执行时间内。

例如:

from numba import jit
import numpy as np
import time

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

@jit(nopython=True)
def go_fast(a): # Function is compiled and runs in machine code
    trace = 0.0
    for i in range(a.shape[0]):
        trace += np.tanh(a[i, i])
    return a + trace

# DO NOT REPORT THIS... COMPILATION TIME IS INCLUDED IN THE EXECUTION TIME!
start = time.perf_counter()
go_fast(x)
end = time.perf_counter()
print("Elapsed (with compilation) = {}s".format((end - start)))

# NOW THE FUNCTION IS COMPILED, RE-TIME IT EXECUTING FROM CACHE
start = time.perf_counter()
go_fast(x)
end = time.perf_counter()
print("Elapsed (after compilation) = {}s".format((end - start)))

例如,这将打印:

Elapsed (with compilation) = 0.33030009269714355s
Elapsed (after compilation) = 6.67572021484375e-06s

衡量 Numba JIT 对代码影响的良好方法是使用 timeit 模块函数计时执行;这些函数测量多次执行的迭代,因此可以适应第一次执行中的编译时间。

顺便提一下,如果编译时间是问题,Numba JIT 支持 编译函数的磁盘缓存 ,并且还有一个 提前编译 模式。

它有多快?

假设 Numba 可以在 nopython 模式下运行,或者至少编译一些循环,它将针对您的特定 CPU 进行编译。加速效果因应用程序而异,但可以达到一到两个数量级。Numba 有一个 性能指南,涵盖了获得额外性能的常见选项。

Numba 是如何工作的?

Numba 读取装饰函数的 Python 字节码,并结合函数输入参数的类型信息。它分析和优化你的代码,最后使用 LLVM 编译器库生成适合你 CPU 能力的机器码版本的函数。这个编译版本在每次调用你的函数时都会被使用。

其他有趣的事情:

Numba 有很多装饰器,我们见过 @jit,但还有:

  • @njit - 这是 @jit(nopython=True) 的别名,因为它非常常用!

  • @vectorize - 生成 NumPy ufunc``(支持所有 ``ufunc 方法)。文档在此

  • @guvectorize - 生成 NumPy 广义 ufunc文档在此

  • @stencil - 声明一个函数作为类似模板操作的内核。文档在这里

  • @jitclass - 用于jit感知的类。文档在这里

  • @cfunc - 声明一个用于作为本地回调函数(从C/C++等调用)。文档在这里

  • @overload - 为在 nopython 模式下使用的函数注册您自己的实现,例如 @overload(scipy.special.j0)文档在这里

一些装饰器中可用的额外选项:

ctypes/cffi/cython 互操作性:

  • cffi - 在 nopython 模式下支持调用 CFFI 函数。

  • ctypes - 在 nopython 模式下支持调用 ctypes 包装的函数。

  • Cython 导出的函数 是可调用的

GPU 目标:

Numba 可以针对 Nvidia CUDA GPU。你可以在纯 Python 中编写内核,并让 Numba 处理计算和数据移动(或明确地执行此操作)。点击查看 Numba 关于 CUDA 的文档。