Numba 包使用简要总结

Numba 是一个开源 JIT 编译工具,功能是将 Python 和 Numpy 代码快速转换为机器码。通过llvmlite Python包,使用LLVM将一部分Python和NumPy (需要注意 Pandas 不能被 Numba 理解,直接对 Pandas 处理导致的结果是计算成本会增加) 转换为快速机器代码。它提供了一系列选项,用于并行化 CPU 和 GPU 的 Python 代码,通常只需进行少量代码更改。

Numpa 的装饰器类型

除了 @jit 外还有其他类型的装饰器:

  1. @njit 这是@jit(nopython = True)的别名,可用于替换
  2. @vectorize 生成NumPy ufunc(支持所有ufunc方法 Universal functions(ufunc))。文件在这里。Numba 能将 Python 函数转换为 ufunc,这样可以实现 Numpy array 的 C 语言方式执行
  3. @guvectorize 生成 NumPy广义ufunc。和 @vectorize 差异在于能够处理任意数量的数组
  4. @stencil 将函数声明为类似模板的操作的内核。是一种通用的计算模式,以每个元素均可以被固定的模式更新数据
  5. @jitclass 用于处理类的编译
  6. @cfunc 申明调用本地的 C/C++ 库。它的使用和 @jit 相似,但是需要强制性给定 signature(可能是需要强制说明数据类型?!需要确认)
  7. @overload-注册自己的函数实现以在nopython模式下使用,例如@overload(scipy.special.j0)

@jit 装饰器

@jit 编译有两种方式:Lazy Compilation 和 Eager Compilation。前者的编译,是在首次执行时才会编译;后者的使用是通过函数 signature 来实现,如果是需要进行详细的精度控制这是一种良好的方式——特别需要注意在某些情况下是否申明了返回值类型会得到不同的结构:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
# Eager 编译,没有申明返回值类型
In [2]: @nb.jit((nb.int32, nb.int32))
...: def f(x, y):
...: return x + y
...:

In [3]: f(1, 2)
Out[3]: 3

In [4]: f(2 **31, 2**31+1)
Out[4]: -4294967295

# Eager 编译,申明了返回值类型,超过精度之后会被抛弃
In [5]: @nb.jit(nb.int32(nb.int32, nb.int32))
...: def f(x, y):
...: return x + y
...:

In [6]: f(2**31, 2**31+1)
Out[6]: 1

# 直接用 Python 进行计算
In [7]: 2**31 + 2**32+1
Out[7]: 6442450945

行内调用其他函数

需要调用其他编译的函数时,建议在调用的函数/类中使用 Numba 编译,否则可能会导致结果很慢:

1
2
3
4
5
6
7
8
@jit
def square(x):
return x ** 2

@jit
def hypot(x, y):
# 调用了编译的 square 函数
return math.sqrt(square(x) + square(y))

nopython 参数

Numpa 的使用一般是通过装饰器的方式,主要有两种模式:nopython 模式和 object 模式。

1
2
3
4
5
6
7
8
9
10
11
12
13
from numba import jit
import numpy as np

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

@jit(nopython=True) # Set "nopython" mode for best performance, equivalent to @njit
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))

nopython 的模式是进行完全编译,而不调用 Python 解释器的方式。这种模式是推荐的最佳实践方式。如果是没有显性的说明 nopython 那么会变为 object 模式,这种模式下是会检查代码中哪些部分可以编译为机器码,其他不能编译的会使用 python 解释器来处理。例如:

1
2
3
4
5
6
7
8
9
10
11
12
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))

上面的代码中因为 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 return None when called from Python)
  • intp and uintp are pointer-sized integers (signed and unsigned, respectively)
  • intc and uintc are equivalent to C int and unsigned int integer types
  • int8, uint8, int16, uint16, int32, uint32, int64, uint64 are fixed-width integers of the corresponding bit width (signed and unsigned)
  • float32 and float64 are single- and double-precision floating-point numbers, respectively
  • complex64 and complex128 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 or int8[:,:] for a two-dimensional array of 8-bit integers.

@genrated_jit 装饰器

该方式的编译是用于解决一个函数执行不同类型输入,实现保留 JIT 函数执行效率的同时,支持选择编译时(compile-time) 的不同类型选择。例如需要根据不同的值返回是否缺失的情况:

  • 输入数据是浮点数类型的却是 NaN
  • Numpy 的时间日期类型缺失
  • 以及不满足缺失条件的数据
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
import numpy as np

from numba import genrate_jit, types

@generated_jit(nopython=True)
def is_missing(x):
"""
Return True if the value is missing, False otherwise.
"""
if isinstance(x, types.Float):
return lambda x: np.isnan(x)
elif isinstance(x, (types.NPDatetime, types.NPTimedelta)):
# The corresponding Not-a-Time value
missing = x('NaT')
return lambda x: x == missing
else:
return lambda x: False

在使用该方式编译的过程中,需要注意:

  1. 装饰器函数需要调用的是 numbatypes 参数
  2. 这类型的装饰器函数,并没有直接返回计算结果,而是返回的一个可调用对象
  3. 预先计算的数据在编译执行时,是可以被重用的
  4. 函数定义使用与修饰后的函数中的参数相同的名称,这是确保通过名称传递参数按预期工作所必需的

可使用参数

该编译模式,使用参数包括 nopythoncache 选项

@vectorize 和 @guvectorize装饰器

这两种装饰器分别处理的是两种类型的数据:

  1. 标量数据计算,这类型是使用 universal functions (或 ufuncs),是直接使用 @vectorize 装饰器
  2. 高维度数据计算,这类型需要使用 generalized universal functions(或 gufuncs),是使用 @guvectorize 装饰器

@vectorize 应用

需要注意在文档中描述的标量数据计算,实际上是解决的是一维 array 数据:

1
2
3
4
5
6
# 主要可以看作是解决了类似下面的直接用 python 代码写出可以逐个处理 array 的元素
import numpy as np
import math

def test(a, b):
return math.sin(a) * math.exp(b)

上面的例子中,如果传入一个 array 数据就会报类型错误。numba 的函数则可以用于解决该问题。

该装饰器也包括了两种运算模式——Eager 编译和 Lazy 编译。在使用 Eager 方式时,可以申明使用多种数据类型,但是要注意 signature 顺序——需要将最不确定的类型放在最后,例如:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
from numba import vectorize, float64, float32, int32, int64

@vectorize([int32(int32, int32),
int64(int64, int64),
float32(float32, float32),
float64(float64, float64)])
def test(x, y):
return x + y

# 计算的结果会根据相应的输入类型匹配条件返回相应的类型
a = np.arange(12)
test(a, a) # 返回的元素类型是 int

a = np.linspace(0, 1, 4)
test(a, a) # 返回的元素类型是 float

# 但是在处理其他类型数据时回报类型错误
a = np.linspace(0, 1+1j, 12)
test(a, a) # 回报类型错误

在实际应用中,可以模仿出 Numpy 的广播方式,例如:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
  >>> a = np.arange(12).reshape(3, 4)
>>> a
array([[ 0, 1, 2, 3],
[ 4, 5, 6, 7],
[ 8, 9, 10, 11]])
>>> test.reduce(a, axis=0)
array([12, 15, 18, 21])
>>> test.reduce(a, axis=1)
array([ 6, 22, 38])
>>> test.accumulate(a)
array([[ 0, 1, 2, 3],
[ 4, 6, 8, 10],
[12, 15, 18, 21]])
>>> test.accumulate(a, axis=1)
array([[ 0, 1, 3, 6],
[ 4, 9, 15, 22],
[ 8, 17, 27, 38]])

target 参数

是用于选择函数以什么样形式计算,cpu 是以单核 cpuparallel 是以多核 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

注意

  1. 使用 numba 的函数,不建议传入 list 会报错 ——“Reflected list” is being deprecated when there is no reflection?
    至于相关的原因,在 核心开发者有解释 。虽然可以通过其他调整属性 numba 的 List,但是没有传入 tuple 的效率高
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
import numba as nb
from numba.typed import List

@nb.njit
def euclidean(x:tuple, y:tuple)->float:
"""
直接计算欧式距离
$\sqrt{\displaystyle{\sum_{i=1}^n(x_i-y_i)^2}}$
"""
if len(x) != len(y):
raise ValueError("Differente Length")

result = 0

for x_i, y_i in zip(x, y):
result += (x_i - y_i) ** 2
return result
  1. Numba 多线程数量设置

    numba的多线程的数量通过全局变量来设置:

    1
    2
    import numba
    numba.config.NUMBA_NUM_THREADS=8
  2. 使用 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 位类型非常重要
  3. 允许直接在 GPU 中运算的 python 表达式

    • if/elif/else
    • while for 循环
    • 基本数学运算
    • mathcmath 模块中的某些函数
    • 元组,Tuple

    详情参考 the Numba manual

作者

ZenRay

发布于

2021-03-12

更新于

2021-03-12

许可协议

CC BY-NC-SA 4.0