Skip to content

1.6。创建 Numpy 通用函数

原文： http://numba.pydata.org/numba-doc/latest/user/vectorize.html

1.6.1。 @vectorize装饰器

Numba 的 vectorize 允许 Python 函数将标量输入参数用作 NumPy ufuncs 。创建传统的 NumPy ufunc 并不是最直接的过程，而是涉及编写一些 C 代码。 Numba 让这很容易。使用 vectorize() 装饰器，Numba 可以将纯 Python 函数编译为一个 ufunc，它在 NumPy 阵列上运行的速度与用 C 编写的传统 ufunc 一样快。

使用 vectorize() ，可以将函数编写为通过输入标量而不是数组进行操作。 Numba 将生成周围循环（或 _ 内核 _），允许对实际输入进行有效迭代。

vectorize() 装饰器有两种操作模式：

  • 渴望或装饰时间，编译：如果您将一个或多个类型签名传递给装饰器，您将构建 Numpy 通用函数（ufunc）。本小节的其余部分描述了使用装饰时编译构建 ufunc。
  • 懒惰或调用时编译：当没有给出任何签名时，装饰器会给你一个 Numba 动态通用函数（ DUFunc ），它在使用以前不支持的输入类型调用时动态编译新内核。后面的小节“动态通用函数”更深入地描述了这种模式。

如上所述，如果您将签名列表传递给 vectorize() 装饰器，您的函数将被编译为 Numpy ufunc。在基本情况下，只会传递一个签名：

from numba import vectorize, float64

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

如果您传递了多个签名，请注意必须在最不具体的签名之前传递大多数特定签名（例如，在双精度浮点数之前单精度浮点数），否则基于类型的分派将无法按预期工作：

@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''

你可能会问自己，“为什么我要经历这个而不是使用 @jit 装饰器编译一个简单的迭代循环？”。答案是 NumPy ufuncs 会自动获得其他功能，如缩小，累积或广播。使用上面的例子：

>>> 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]])

也可以看看

ufuncs 的标准功能（NumPy 文档）。

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

目标 描述
中央处理器 单线程 CPU
平行 多核 CPU
CUDA CUDA GPU 注意这会创建一个类似 _ufunc 的 _ 对象。有关详细信息，请参阅 CUDA ufunc 的文档。

一般准则是为不同的数据大小和算法选择不同的目标。 “cpu”目标适用于小数据大小（约小于 1KB）和低计算强度算法。它具有最少的开销。 “并行”目标适用于中等数据大小（大约小于 1MB）。线程增加了一点延迟。 “cuda”目标适用于大数据量（大约 1MB）和高计算强度算法。向 GPU 传输内存和从 GPU 传输内存会增加大量开销。

1.6.2。 @guvectorize装饰器

虽然 vectorize() 允许你一次写一个元素的 ufunc，但 guvectorize() 装饰器更进一步，并允许你编写可以工作的 ufuncs 在输入数组的任意数量的元素上，并获取和返回不同维度的数组。典型的例子是运行中值或卷积滤波器。

与 vectorize() 函数相反， guvectorize() 函数不返回其结果值：它们将其作为数组参数，必须由函数填充。这是因为数组实际上是由 NumPy 的调度机制分配的，调用机制调用 Numba 生成的代码。

这是一个非常简单的例子：

@guvectorize([(int64[:], int64, int64[:])], '(n),()->(n)')
def g(x, y, res):
    for i in range(x.shape[0]):
        res[i] = x[i] + y

底层的 Python 函数只是将一个给定的标量（y）添加到一维数组的所有元素中。宣言更有意思。那里有两件事：

  • 输入和输出 _ 布局 _ 的声明，符号形式：(n),()->(n)告诉 NumPy 该函数采用 n 元素一维数组，一个标量（用符号表示为空元组()）并返回 n 元素一维数组;
  • @vectorize中支持的具体 _ 签名 _ 列表;这里我们只支持int64数组。

注意

1D 数组类型也可以接收标量参数（形状为()的参数）。在上面的例子中，第二个参数也可以声明为int64[:]。在这种情况下，该值必须由y[0]读取。

我们现在可以通过一个简单的例子来检查已编译的 ufunc 的作用：

>>> a = np.arange(5)
>>> a
array([0, 1, 2, 3, 4])
>>> g(a, 2)
array([2, 3, 4, 5, 6])

好处是 NumPy 将根据其形状自动调度更复杂的输入：

>>> a = np.arange(6).reshape(2, 3)
>>> a
array([[0, 1, 2],
       [3, 4, 5]])
>>> g(a, 10)
array([[10, 11, 12],
       [13, 14, 15]])
>>> g(a, np.array([10, 20]))
array([[10, 11, 12],
       [23, 24, 25]])

注意

vectorize() 和 guvectorize() 都支持传递nopython=True ，如同@jit 装饰器。使用它来确保生成的代码不会回退到对象模式。

1.6.3。动态通用功能

如上所述，如果您没有将任何签名传递给 vectorize() 装饰器，您的 Python 函数将用于构建动态通用函数，或 DUFunc 。例如：

from numba import vectorize

@vectorize
def f(x, y):
    return x * y

结果f()是 DUFunc 实例，以没有支持的输入类型开头。在调用f()时，只要传递以前不支持的输入类型，Numba 就会生成新的内核。鉴于上面的示例，以下一组解释器交互说明了动态编译的工作原理：

>>> f
<numba._DUFunc 'f'>
>>> f.ufunc
<ufunc 'f'>
>>> f.ufunc.types
[]

上面的例子显示 DUFunc 实例不是 ufunc。而不是子类 ufunc， DUFunc 实例通过保持 ufunc 成员，然后将 ufunc 属性读取和方法调用委托给此成员（也称为类型聚合）来工作。当我们查看 ufunc 支持的初始类型时，我们可以验证没有。

我们试着打电话给f()：

>>> f(3,4)
12
>>> f.types   # shorthand for f.ufunc.types
['ll->l']

如果这是一个普通的 Numpy ufunc，我们会看到一个异常抱怨 ufunc 无法处理输入类型。当我们用整数参数调用f()时，我们不仅会收到答案，而且我们可以验证 Numba 是否创建了支持 C long整数的循环。

我们可以通过使用不同的输入调用f()来添加其他循环：

>>> f(1.,2.)
2.0
>>> f.types
['ll->l', 'dd->d']

我们现在可以验证 Numba 是否为处理浮点输入添加了第二个循环"dd->d"。

如果我们将输入类型混合到f()，我们可以验证 Numpy ufunc 强制转换规则是否仍然有效：

>>> f(1,2.)
2.0
>>> f.types
['ll->l', 'dd->d']

此示例演示了使用混合类型调用f()会导致 Numpy 选择浮点循环，并将整数参数转换为浮点值。因此，Numba 没有创建一个特殊的"dl->d"内核。

这 DUFunc 行为导致我们得到类似于上面“ @vectorize 装饰器”小节中给出的警告的点，但是在装饰器中没有签名声明顺序，呼叫顺序很重要。如果我们首先传入浮点参数，那么任何带有整数参数的调用都将被转换为双精度浮点值。例如：

>>> @vectorize
... def g(a, b): return a / b
...
>>> g(2.,3.)
0.66666666666666663
>>> g(2,3)
0.66666666666666663
>>> g.types
['dd->d']

如果您需要对各种类型签名的精确支持，您应该在 vectorize() 装饰器中指定它们，而不是依赖于动态编译。

回到顶部