Skip to content

1.1. Numba 的 ~ 5 分钟指南

原文: http://numba.pydata.org/numba-doc/latest/user/5minguide.html

Numba 是 Python 的即时(JIT)编译器,它最适用于使用 NumPy 数组、函数和循环的代码。使用 Numba 的最常用方法是使用其装饰器集合,通过作用于您的函数来指示 Numba 编译它们。当调用 Numba 修饰函数时,它被编译为机器代码“即时”(JIT)执行,从而使您的全部或部分代码随后可以以本机机器代码速度运行!

Numba 可直接运行于下列环境:

  • 操作系统:Windows(32 位和 64 位),OSX 和 Linux(32 位和 64 位)
  • 架构:x86,x86_64,ppc64le。在 armv7l,armv8l(aarch64)上进行实验。
  • GPU:Nvidia CUDA。 在 AMD ROC 上的应用是实验性质的。
  • CPython
  • NumPy 1.10 - 最新的

1.1.1. 如何获取Numba?

Numba 是 Anaconda Python 发行版的一个 conda 软件包:

$ conda install numba

Numba 还有wheels(.whl)可供选择:

$ pip install numba

Numba 也可以从源码 (校对注:md超链接需要整体修改) 编译,但我们不建议首次使用 Numba 的用户使用这种方法。

Numba 通常用作核心包,因此其依赖性被尽可能减小,但可以额外安装下列包以提供附加功能:

  • scipy - 支持编译 numpy.linalg 的函数。
  • colorama - 支持回溯/错误消息中的颜色突出显示。
  • pyyaml - 启用 Numba 配置的 YAML 配置文件。
  • icc_rt - 允许使用 Intel SVML(高性能短矢量数学库,仅限 x86_64 架构)。安装说明在性能提示 (校对注:md超链接需要整体修改) 中。

1.1.2. Numba 会对我的代码有用吗?

这取决于你的代码是什么样的,如果你的代码是以数字为导向的(做很多数学运算),使用 NumPy 和/或有很多循环,那么 Numba 通常是一个不错的选择。在以下的例子中,我们将应用最基本的 Numba 的 JIT 装饰器@jit来尝试加速某些函数,以展示哪些函数会被改善,哪些函数不能被改善。

Numba 可以改善的代码看起来像这样:

from numba import jit
import numpy as np

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

@jit(nopython=True) # 使用 "nopython" 模式以提供更好的性能, 等价于 @njit
def go_fast(a): # 当函数第一次被调用时会被编译为机器码
    trace = 0
    for i in range(a.shape[0]):   # Numba 偏好循环
        trace += np.tanh(a[i, i]) # Numba 偏好 NumPy 函数
    return a + trace              # Numba 偏好 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): # 这个函数无法受益于 Numba jit
    df = pd.DataFrame.from_dict(a) # Numba 不识别 pd.DataFrame
    df += 1                        # Numba 不明白这是什么
    return df.cov()                # 或是这个

print(use_pandas(x))

请注意,Numba 不识别 Pandas,因此 Numba 只是通过解释器运行此代码,却增加了 Numba 内部开销的成本!

1.1.3. 什么是nopython模式?

Numba @jit装饰器基本上以两种编译模式运行,nopython模式和object模式。在上面的go_fast示例中,在@jit装饰器中设置了nopython=True,这是指示 Numba 在nopython模式下运行。 nopython编译模式的行为本质上是编译装饰函数,以便它完全运行而不需要 Python 解释器的参与。这是使用 Numba jit装饰器的推荐和最佳实践方式,因为它可以获得最佳性能。

如果nopython模式下的编译失败,Numba 可以使用object mode进行编译,如果未设置nopython=True,这是@jit装饰器的降级模式(如上面的use_pandas示例所示)。在这种模式下,Numba 将识别它可以编译的循环,并将它们编译成在机器代码中运行的函数,并且它将运行解释器中的其余代码。为获得最佳性能,请避免使用此模式!

1.1.4. 如何衡量 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
    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.time()
go_fast(x)
end = time.time()
print("Elapsed (with compilation) = %s" % (end - start))

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

如此会打印(举例):

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

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

另外,如果编译时间是一个问题,Numba JIT 支持已编译函数的磁盘缓存 (校对注:md超链接需要整体修改) ,并且还具有 Ahead-Of-Time (校对注:md超链接需要整体修改) 编译模式。

1.1.5. Numba有多快?

假设 Numba 可以在nopython模式下运行,或者至少编译一些循环,它将针对您的特定 CPU 进行编译。加速因应用而异,但可以是一到两个数量级。 Numba 有一个性能指南 (校对注:md超链接需要整体修改) ,涵盖了获得额外性能的常用选项。

1.1.6. Numba 如何工作?

Numba 读取被装饰函数的 Python 字节码,并将其与有关函数输入参数类型的信息相结合。Numba会分析并优化您的代码,最后使用 LLVM 编译器库生成函数的机器代码版本,这将针对您的 CPU 功能量身定制。此后每次调用函数时都会使用此编译版本。

1.1.7. 其他有意义的东西:

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

  • @njit - 这是@jit(nopython=True)的别名,因为它非常常用!
  • @vectorize - 产生 NumPy ufunc(支持所有ufunc方法)。 文件在这里 (校对注:md超链接需要整体修改) 。
  • @guvectorize - 产生 NumPy 广义ufunc文件在这里 (校对注:md超链接需要整体修改) 。
  • @stencil - 将函数声明为类似模板操作的内核。 文件在这里 (校对注:md超链接需要整体修改) 。
  • @jitclass - 用于 jit 感知类。 文件在这里 (校对注:md超链接需要整体修改) 。
  • @cfunc - 声明一个用作本机回调的函数(从 C / C ++等调用)。 文件在这里 (校对注:md超链接需要整体修改) 。
  • @overload - 注册您自己的函数实现,以便在 nopython 模式下使用,例如: @overload(scipy.special.j0)文件在这里 (校对注:md超链接需要整体修改) 。

一些装饰器提供额外选项:

  • parallel = True - 开启(校对注:md超链接需要整体修改)函数的自动并行化(校对注:md超链接需要整体修改)。
  • fastmath = True - 为该函数启用 fast-math (校对注:md超链接需要整体修改) 表现。

ctypes / cffi / cython 互用性:

  • cffi - nopython模式支持调用 CFFI (校对注:md超链接需要整体修改)功能。
  • ctypes - nopython模式支持调用 ctypes (校对注:md超链接需要整体修改) 包装函数。
  • Cython 导出函数可调用(校对注:md超链接需要整体修改)。

1.1.7.1. 用于 GPU 目的:

Numba 可以面向 Nvidia CUDAAMD ROC(实验性的) GPU 进行程序设计。您可以用纯 Python 编写内核,让 Numba 处理计算和数据移动(或显式执行此操作)。参见 CUDA(校对注:md超链接需要整体修改) 或 ROC (校对注:md超链接需要整体修改) 的 Numba 文档。



回到顶部