Numba 包使用简要总结
Numba 是一个开源 JIT 编译工具,功能是将 Python 和 Numpy 代码快速转换为机器码。通过llvmlite Python包,使用LLVM将一部分Python和NumPy (需要注意 Pandas 不能被 Numba 理解,直接对 Pandas 处理导致的结果是计算成本会增加) 转换为快速机器代码。它提供了一系列选项,用于并行化 CPU 和 GPU 的 Python 代码,通常只需进行少量代码更改。
Numpa 的装饰器类型
除了 @jit
外还有其他类型的装饰器:
- @njit 这是@jit(nopython = True)的别名,可用于替换
- @vectorize 生成NumPy ufunc(支持所有ufunc方法 Universal functions(ufunc))。文件在这里。Numba 能将 Python 函数转换为 ufunc,这样可以实现 Numpy array 的 C 语言方式执行
- @guvectorize 生成 NumPy广义ufunc。和 @vectorize 差异在于能够处理任意数量的数组
- @stencil 将函数声明为类似模板的操作的内核。是一种通用的计算模式,以每个元素均可以被固定的模式更新数据
- @jitclass 用于处理类的编译
- @cfunc 申明调用本地的 C/C++ 库。它的使用和 @jit 相似,但是需要强制性给定 signature(可能是需要强制说明数据类型?!需要确认)
- @overload-注册自己的函数实现以在nopython模式下使用,例如@overload(scipy.special.j0)
@jit 装饰器
@jit
编译有两种方式:Lazy Compilation 和 Eager Compilation。前者的编译,是在首次执行时才会编译;后者的使用是通过函数 signature 来实现,如果是需要进行详细的精度控制这是一种良好的方式——特别需要注意在某些情况下是否申明了返回值类型会得到不同的结构:
1 | # Eager 编译,没有申明返回值类型 |
行内调用其他函数
需要调用其他编译的函数时,建议在调用的函数/类中使用 Numba 编译,否则可能会导致结果很慢:
1 |
|
nopython 参数
Numpa 的使用一般是通过装饰器的方式,主要有两种模式:nopython
模式和 object
模式。
1 | from numba import jit |
nopython
的模式是进行完全编译,而不调用 Python 解释器的方式。这种模式是推荐的最佳实践方式。如果是没有显性的说明 nopython
那么会变为 object
模式,这种模式下是会检查代码中哪些部分可以编译为机器码,其他不能编译的会使用 python 解释器来处理。例如:
1 | from numba import jit |
上面的代码中因为 Numba 对 pandas 处理效果不好,因此可以通过关闭 nopython
模式让可以通过 python 解释器执行的,通过 python 解释器执行。
此外需要注意,使用 numa 会需要将代码进行编译,这在首次执行时会消耗一定时间用于编译。因此如果是需要评估代码的时间消耗,可以在 IPython 中使用 %timeit 方式进行评估
nogil 参数
nogil
参数是用于控制是否还需要维持使用 Python 的 GIL。释放 GIL 可以充分利用多核处理器计算以及多线程计算。如果是在对象模式下(即代码处理的数据全是以 Python 对象或者使用 Python C 的 API 处理这些对象的模式),也不可能释放 GIL。虽然释放了 GIL 带来了效率,但是注意会带来多线程编程中的隐患。
cache 参数
cache
参数用于控制函数编译结果是否需要写入缓存。为了避免每次调用Python程序时都要进行编译,可以指示Numba将函数编译的结果写入基于文件的缓存中
parallel 参数
parallel = True
是用于允许自动化执行并行函数运算,fastmath = True
作用是可以进行不安全的浮点运算,相对来说准确率要低一些。
Signature 申明
显性的 signature 类型,包括:
void
is the return type of functions returning nothing (which actually returnNone
when called from Python)intp
anduintp
are pointer-sized integers (signed and unsigned, respectively)intc
anduintc
are equivalent to Cint
andunsigned int
integer typesint8
,uint8
,int16
,uint16
,int32
,uint32
,int64
,uint64
are fixed-width integers of the corresponding bit width (signed and unsigned)float32
andfloat64
are single- and double-precision floating-point numbers, respectivelycomplex64
andcomplex128
are single- and double-precision complex numbers, respectively- array types can be specified by indexing any numeric type, e.g.
float32[:]
for a one-dimensional single-precision array orint8[:,:]
for a two-dimensional array of 8-bit integers.
@genrated_jit 装饰器
该方式的编译是用于解决一个函数执行不同类型输入,实现保留 JIT 函数执行效率的同时,支持选择编译时(compile-time) 的不同类型选择。例如需要根据不同的值返回是否缺失的情况:
- 输入数据是浮点数类型的却是
NaN
- Numpy 的时间日期类型缺失
- 以及不满足缺失条件的数据
1 | import numpy as np |
在使用该方式编译的过程中,需要注意:
- 装饰器函数需要调用的是
numba
的types
参数 - 这类型的装饰器函数,并没有直接返回计算结果,而是返回的一个可调用对象
- 预先计算的数据在编译执行时,是可以被重用的
- 函数定义使用与修饰后的函数中的参数相同的名称,这是确保通过名称传递参数按预期工作所必需的
可使用参数
该编译模式,使用参数包括 nopython
和 cache
选项
@vectorize 和 @guvectorize装饰器
这两种装饰器分别处理的是两种类型的数据:
- 标量数据计算,这类型是使用
universal functions
(或ufuncs
),是直接使用@vectorize
装饰器 - 高维度数据计算,这类型需要使用
generalized universal functions
(或gufuncs
),是使用@guvectorize
装饰器
@vectorize 应用
需要注意在文档中描述的标量数据计算,实际上是解决的是一维 array
数据:
1 | # 主要可以看作是解决了类似下面的直接用 python 代码写出可以逐个处理 array 的元素 |
上面的例子中,如果传入一个 array
数据就会报类型错误。numba
的函数则可以用于解决该问题。
该装饰器也包括了两种运算模式——Eager 编译和 Lazy 编译。在使用 Eager 方式时,可以申明使用多种数据类型,但是要注意 signature 顺序——需要将最不确定的类型放在最后,例如:
1 | from numba import vectorize, float64, float32, int32, int64 |
在实际应用中,可以模仿出 Numpy
的广播方式,例如:
1 | >>> a = np.arange(12).reshape(3, 4) |
target 参数
是用于选择函数以什么样形式计算,cpu
是以单核 cpu
,parallel
是以多核 cpu
的方式计算,cuda
是以 CUDA GPU
方式计算。选择什么样的方式计算,可以参考数据量:
- 小数据量数据,小于
1KB
的数据和低计算要求的算法,可以使用cpu
- 中等数据量数据,近似
1MB
的数据可以使用parallel
- 超过
1MB
的数据和高频计算的算法,推荐使用gpu
计算
cache 参数
也是用于选择是否需要将函数缓存
@gpuvectorize 应用
主要是解决任意维度的输入数据计算,
@vectorize 和 @guvectorize 尚未弄清楚
http://numba.pydata.org/numba-doc/latest/user/vectorize.html#guvectorize
http://numba.pydata.org/numba-doc/latest/user/vectorize.html#vectorize
注意
- 使用 numba 的函数,不建议传入 list 会报错 ——“Reflected list” is being deprecated when there is no reflection?
至于相关的原因,在 核心开发者有解释 。虽然可以通过其他调整属性 numba 的 List,但是没有传入 tuple 的效率高
1 | import numba as nb |

Numba 多线程数量设置
numba的多线程的数量通过全局变量来设置:
1
2import numba
numba.config.NUMBA_NUM_THREADS=8使用 GPU 的注意事项,使用 GPU 并不一定会提高运算效率,其原因是
- 输入数据量太小,GPU通过并行性实现性能,同时处理数千个值。需要更大的阵列才能使GPU繁忙
- 计算太简单了,与在CPU上调用函数相比,将计算发送到GPU涉及大量开销。如果计算没有涉及足够的数学运算(通常称为“算术强度(arithmetic intensity)”),则GPU将花费大部分时间等待数据移动
- 数据复制,到GPU或从GPU复制数据,虽然对于单个函数运行包括复制数据到时间。但通常希望依次运行多个GPU操作时,将数据发送到 GPU 并保留在那里直到完成所有处理才有意义
- 数据类型长度一定超过数据需求,使用32位和64位数据类型的标量运行速度在CPU上基本相同,但是64位数据类型在 GPU 上的性能成本却很高。 64位浮点数的基本算术运算速度比32位浮点数慢2倍(Pascal架构Tesla)到24倍(Maxwell架构 GeForce)。 创建数组时,NumPy 默认为64位数据类型,因此设置 dtype属性或在需要时使用
ndarray.astype()
方法选择 32 位类型非常重要
允许直接在 GPU 中运算的 python 表达式
if
/elif
/else
while
for
循环- 基本数学运算
math
和cmath
模块中的某些函数- 元组,Tuple
详情参考 the Numba manual