numba 高级用法

 

numba是一个用于编译Python数组和数值计算函数的编译器，这个编译器能够大幅提高直接使用Python编写的函数的运算速度。

numba使用LLVM编译器架构将纯Python代码生成优化过的机器码，通过一些添加简单的注解，将面向数组和使用大量数学的python代码优化到与c，c++和Fortran类似的性能，而无需改变Python的解释器。

Numba的主要特性：

• 动态代码生成 （在用户偏爱的导入期和运行期）
• 为CPU（默认）和GPU硬件生成原生的代码
• 集成Python的科学软件栈（Numpy）

下面是使用Numba优化的函数方法，将Numpy数组作为参数：

import numba
@numba.jit
def sum2d(arr):
    M, N = arr.shape
    result = 0.0
    for i in range(M):
        for j in range(N):
            result += arr[i,j]
    return result

如果你对此不是太感兴趣，或者对于其他的加速方案已经很熟悉，可以到此为止，只需要了解加上jit装饰器就可以实现了。

使用jit

使用jit的好处就在于让numba来决定什么时候以及怎么做优化。

from numba import jit

@jit
def f(x, y):
    # A somewhat trivial example
    return x + y

比如这段代码，计算将延期到第一次函数执行，numba将在调用期间推断参数类型，然后基于这个信息生成优化后的代码。numba也能够基于输入的类型编译生成特定的代码。例如，对于上面的代码，传入整数和复数作为参数将会生成不同的代码：

>>>f(1,2)
3
>>>f(1j,2)
(2+1j)

我们也可以加上所期望的函数签名：

from numba import jit, int32

@jit(int32(int32, int32))
def f(x, y):
    # A somewhat trivial example
    return x + y

int32(int32, int32) 是函数签名，这样，相应的特性将会被@jit装饰器编译，然后，编译器将控制类型选择，并不允许其他特性（即其他类型的参数输入，如float）

Numba编译的函数可以调用其他编译函数。 函数调用甚至可以在本机代码中内联，具体取决于优化器的启发式。 例如：

@jit
def square(x):
    return x ** 2

@jit
def hypot(x, y):
    return math.sqrt(square(x) + square(y))

@jit装饰器必须添加到任何库函数，否则numba可能生成速度更慢的代码。

签名规范

Explicit @jit signatures can use a number of types. Here are some common ones:

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.

编译选项

numba有两种编译模式：nopython模式和object模式。前者能够生成更快的代码，但是有一些限制可能迫使numba退为后者。想要避免退为后者，而且抛出异常，可以传递nopython=True.

@jit(nopython=True)
def f(x, y):
    return x + y

当Numba不需要保持全局线程锁时，如果用户设定nogil=True，当进入这类编译好的函数时，Numba将会释放全局线程锁。

@jit(nogil=True)
def f(x, y):
    return x + y

这样可以利用多核系统，但不能使用的函数是在object模式下编译。

想要避免你调用python程序的编译时间，可以这顶numba保存函数编译结果到一个基于文件的缓存中。可以通过传递cache=True实现。

@jit(cache=True)
def f(x, y):
    return x + y

开启一个实验性质的特性将函数中的这些操作自动并行化。这一特性可以通过传递parallel=True打开，然后必须也要和nopython=True配合起来一起使用。编译器将编译一个版本，并行运行多个原生的线程（没有GIL）

@jit(nopython=True, parallel=True)
def f(x, y):
    return x + y

generated_jit

有时候想要编写一个函数，基于输入的类型实现不同的实现，generated_jit()装饰器允许用户在编译期控制不同的特性的选择。假定想要编写一个函数，基于某些需求，返回所给定的值是否缺失的类型，具体定义如下：

• 对于浮点数，缺失的值为NaN。
• 对于Numpy的datetime64和timedelta64参数，缺失值为NaT
• 其他类型没有定义的缺失值

import numpy as np

from numba import generated_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. 调用装饰器函数是使用Numba的类型作为参数，而不是他们的值。
2. 装饰器函数并不真的计算结果，而是返回一个对于给定类型，可调用的实际定义的函数执行。
3. 可以在编译期预先计算一些数据，使其在编译后执行过程中重用。
4. 函数定义使用和装饰器函数中相同名字的参数，这将确保通过名字传递参数能够如期望的工作。

使用@vectorize 装饰器创建Numpy的 universal 函数

Numba的vectorize允许Python函数将标量输入参数作为Numpy的ufunc使用，将纯Python函数编译成ufunc，使之速度与使用c编写的传统的ufunc函数一样。

vectorize()有两种操作模型：

1. 主动，或者装饰期间编译：如果传递一个或者多个类型签名给装饰器，就将构建Numpy的universal function。后面将介绍使用装饰期间编译ufunc。
2. 被动（惰性），或者调用期间编译：当没有提供任何签名，装饰器将提供一个Numba动态universal function(DUFunc)，当一个未支持的新类型调用时，就动态编译一个新的内核，后面的“动态 universal functions”将详细介绍

如上所描述，如果传递一个签名给vectorizer()装饰器，函数将编译成一个numpy 的ufunc：

from numba import vectorize, float64

@vectorize([float64(float64, float64)])
def f(x, y):
    return x + y

如果想传递多个签名，注意顺序，精度低的在前，高的在后，否则就会出奇怪的问题。例如int32就只能在int64之前。

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

如果给定的类型正确：

>>> a = np.arange(6)
>>> f(a, a)
array([ 0,  2,  4,  6,  8, 10])
>>> a = np.linspace(0, 1, 6)
>>> f(a, a)
array([ 0. ,  0.4,  0.8,  1.2,  1.6,  2. ])

如果提供了不支持的类型：

>>> a = np.linspace(0, 1+1j, 6)
>>> f(a, a)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: ufunc 'ufunc' not supported for the input types, and the inputs could not be safely coerced to any supported types according to the casting rule ''safe''

vectorizer与jit装饰器的差别：numpy的ufunc自动加载其他特性，例如：reduction, accumulation or broadcasting:

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

vectorize() 装饰器支持多个ufunc 目标:

Target	Description
cpu	Single-threaded CPU
parallel	Multi-core CPU
cuda	CUDA GPU

guvectorize装饰器只用了进一步的概念，允许用户编写ufuncs操作输入数组中的任意数量的元素，返回不同纬度的数组。典型的应用是运行求均值或者卷积滤波。

Numba支持通过jitclass装饰器实现对于类的代码生成。可以使用这个装饰器来标注优化，类中的所有方法都被编译成nopython function。

import numpy as np
from numba import jitclass          # import the decorator
from numba import int32, float32    # import the types

spec = [
    ('value', int32),               # a simple scalar field
    ('array', float32[:]),          # an array field
]

@jitclass(spec)
class Bag(object):
    def __init__(self, value):
        self.value = value
        self.array = np.zeros(value, dtype=np.float32)

    @property
    def size(self):
        return self.array.size

    def increment(self, val):
        for i in range(self.size):
            self.array[i] = val
        return self.array

性能建议

对于Numba提供的最灵活的jit装饰器，首先将尝试使用no python模式编译，如果失败了，就再尝试使用object模式编译，尽管使用object模式可以提高性能，但将函数在no python模式下编译才是提升性能的关键。想要直接使用nopython模式，可以直接使用装饰器@njit，这个装饰器与@jit(nopython=True)等价。

@njit
def ident_np(x):
    return np.cos(x) ** 2 + np.sin(x) ** 2

@njit
def ident_loops(x):
    r = np.empty_like(x)
    n = len(x)
    for i in range(n):
        r[i] = np.cos(x[i]) ** 2 + np.sin(x[i]) ** 2
    return r

Function Name	@njit	Execution time
ident_np	No	0.581s
ident_np	Yes	0.659s
ident_loops	No	25.2s
ident_loops	Yes	0.670s

有时候不那么严格的规定数据将会带来性能的提升，此时，恶意使用fastmath关键字参数：

@njit(fastmath=False)
def do_sum(A):
    acc = 0.
    # without fastmath, this loop must accumulate in strict order
    for x in A:
        acc += np.sqrt(x)
    return acc

@njit(fastmath=True)
def do_sum_fast(A):
    acc = 0.
    # with fastmath, the reduction can be vectorized as floating point
    # reassociation is permitted.
    for x in A:
        acc += np.sqrt(x)
    return acc

Function Name	Execution time
do_sum	35.2 ms
do_sum_fast	17.8 ms

Trubleshooting and tips

想要编译什么？

通常建议是编译代码中耗时最长的关键路径，如果有一部分代码耗时很长，但在一些高阶的代码之中，可能就需要重构这些对于性能有更高要求的代码到一个单独的函数中，让numba专注于这些对于性能敏感的代码有以下好处：

1. 避免遇见不支持的特性
2. 减少编译时间
3. 在需要编译的函数外，高阶的代码会更简单

不想要编译什么？

numba编译失败的原因很多，最常见的一个原因就是你写的代码依赖于不支持的Python特性，尤其是nopython模式，可以查看支持的python特性

在numba编译代码之前，先要确定所有使用的变量的类型，这样就能生成你的代码的特定类型的机器码。一个常见的编译失败原因（尤其是nopython模式）就是类型推导失败，numba不能确定代码中所有变量的类型。

例如：参考这个函数：

@jit(nopython=True)
def f(x, y):
    return x + y

如果使用两个数字作为参数：

>>> f(1,2)
    3

如果传入一个元组和一个数字，numba不能得到数字和元组求和的结果，就会触发编译报错：

>>> f(1, (2,))
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
File "<path>/numba/numba/dispatcher.py", line 339, in _compile_for_args
    reraise(type(e), e, None)
File "<path>/numba/numba/six.py", line 658, in reraise
    raise value.with_traceback(tb)
numba.errors.TypingError: Failed at nopython (nopython frontend)
Invalid usage of + with parameters (int64, tuple(int64 x 1))
Known signatures:
* (int64, int64) -> int64
* (int64, uint64) -> int64
* (uint64, int64) -> int64
* (uint64, uint64) -> uint64
* (float32, float32) -> float32
* (float64, float64) -> float64
* (complex64, complex64) -> complex64
* (complex128, complex128) -> complex128
* (uint16,) -> uint64
* (uint8,) -> uint64
* (uint64,) -> uint64
* (uint32,) -> uint64
* (int16,) -> int64
* (int64,) -> int64
* (int8,) -> int64
* (int32,) -> int64
* (float32,) -> float32
* (float64,) -> float64
* (complex64,) -> complex64
* (complex128,) -> complex128
* parameterized
[1] During: typing of intrinsic-call at <stdin> (3)

File "<stdin>", line 3:

错误信息“Invalid usage of + with parameters (int64, tuple(int64 x 1))”可以解释为numba解释器遇到了一个整数和元组中的整数求和，

类型统一问题

另一个编译失败的常见原因是：不能静态的决定返回的类型；返回值的类型仅仅依赖于运行期。这样的事情也是仅仅发生在nopython 模式下。类型统一的概念仅仅只是尝试找到一个类型，两个变量能够使用该类型安全的显示；例如一个64位的浮点数和一个64位的复数可以同时使用128位的复数表示。

以下是一个类型统一错误，这个函数的返回类型是基于x的值在运行期决定的：

In [1]: from numba import jit

In [2]: @jit(nopython=True)
...: def f(x):
...:     if x > 10:
...:         return (1,)
...:     else:
...:         return 1
...:

尝试执行这个函数，就会得到以下的错误：

In [3]: f(10)
TypingError: Failed at nopython (nopython frontend)
Can't unify return type from the following types: tuple(int64 x 1), int64
Return of: IR name '$8.2', type '(int64 x 1)', location:
File "<ipython-input-2-51ef1cc64bea>", line 4:
def f(x):
    <source elided>
    if x > 10:
        return (1,)
        ^
Return of: IR name '$12.2', type 'int64', location:
File "<ipython-input-2-51ef1cc64bea>", line 6:
def f(x):
    <source elided>
    else:
        return 1

错误信息： “Can’t unify return type from the following types: tuple(int64 x 1), int64” 可以理解为： “Numba cannot find a type that can safely represent a 1-tuple of integer and an integer”.

编译的太慢

最常见的编译速度很慢的原因是：nopython模式编译失败，然后尝试使用object模式编译。object模式当前几乎没有提供加速特性，只是提供了一种叫做loop-lifting的优化，这个优化将允许使用nopython模式在内联迭代下编译。

可以在编译好的函数上使用inspect_types()方法来查看函数的类型推导是否成功。例如，对于以下函数：

@jit
def f(a, b):
    s = a + float(b)
    return s

当使用numbers调用时，该函数将和numba一样快速的将数字转换为浮点数：

>>> f(1, 2)
3.0
>>> f.inspect_types()
f (int64, int64)
--------------------------------------------------------------------------------
# --- LINE 7 ---

@jit

# --- LINE 8 ---

def f(a, b):

    # --- LINE 9 ---
    # label 0
    #   a.1 = a  :: int64
    #   del a
    #   b.1 = b  :: int64
    #   del b
    #   $0.2 = global(float: <class 'float'>)  :: Function(<class 'float'>)
    #   $0.4 = call $0.2(b.1, )  :: (int64,) -> float64
    #   del b.1
    #   del $0.2
    #   $0.5 = a.1 + $0.4  :: float64
    #   del a.1
    #   del $0.4
    #   s = $0.5  :: float64
    #   del $0.5

    s = a + float(b)

    # --- LINE 10 ---
    #   $0.7 = cast(value=s)  :: float64
    #   del s
    #   return $0.7

    return s

关闭jit编译

设定NUMBA_DISABLE_JIT 环境变量为 1.

FAQ

Q：能否传递一个函数作为参数？
A：不能，但可以使用闭包来模拟实现，例如：

@jit(nopython=True)
def f(g, x):
    return g(x) + g(-x)

result = f(my_g_function, 1)

可以使用一个工厂函数重构：

def make_f(g):
    # Note: a new f() is compiled each time make_f() is called!
    @jit(nopython=True)
    def f(x):
        return g(x) + g(-x)
    return f

f = make_f(my_g_function)
result = f(1)

Q：对于全局变量修改的问题
A：非常不建议使用全局变量，否则只能使用recompile()函数重新编译，这样还不如重构代码，不使用全局变量。

Q：如何调试jit的函数？
A：可以调用pdb，也可以临时关闭编译环境变量：NUMBA_DISABLE_JIT。

Q：如何增加整数的位宽
A：默认情况下，numba为整形变量生成机器整形位宽。我们可以使用np.int64为相关变量初始化（例如：np.int64(0)而不是0）。

Q：如何知道parallel=True已经工作了？
A：如果parallel=True，设定环境变量NUMBA_WARNING为非0，所装饰的函数转换失败，就显示报警；同样，环境变量：NUMBA_DEBUG_ARRAY_OPT_STAT将展示一些统计结果。