Question

我想使用 Numba 对要评估矩阵每一行的函数进行矢量化处理。与循环遍历行相反，这实际上会将Numpy ufunc应用于矩阵。根据{{3}}：

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

考虑到这一点，我什至无法得到一个玩具例子。以下简单示例尝试计算每一行中的元素总数。

import numba, numpy as np

# Define the row-wise function to be vectorized:
@numba.guvectorize(["void(float64[:],float64)"],"(n)->()")
def f(a,b):
    b = a.sum() 

# Apply the function to an array with five rows:
a = np.arange(10).reshape(5,2)
b = f(a)

我使用了@guvectorize装饰器，因为我希望装饰函数将参数 a 用作矩阵的每一行，矩阵是数组； @vectorize仅接受标量输入。我还编写了签名以采用数组参数并修改标量输出。根据{{3}}，修饰函数不使用return语句。

结果应为b = [1,5,9,13,17]，但我得到了b=[0.,1.,2.,3.,4.]。显然，我缺少了一些东西。我希望能有一些指导，请记住，总和只是一个示例。

Answer 1

b = a.sum()永远无法使用python语法修改b的原始值。

numba通过要求将gufunc的每个参数都为数组来解决此问题-标量只是长度1，您可以将其分配为该长度。因此，您需要将两个参数都作为数组，并且赋值必须使用[]

@numba.guvectorize(["void(float64[:],float64[:])"],"(n)->()")
def f(a,b):
    b[:] = a.sum()
    # or b[0] = a.sum()

f(a)
Out[246]: array([ 1.,  5.,  9., 13., 17.])

Answer 2

@chrisb 上面有一个很好的答案。这个答案应该为那些刚接触矢量化的人添加一些说明。

在向量化方面（在 numpy 和 numba 中），您传递输入向量。

例如：

import numpy as np

a=[1,2]
b=[3,4]

@np.vectorize
def f(x_1,x_2):
    return x_1+x_2

print(f(a,b))
#-> [4,6]

在 numba 中，传统上您需要将输入类型传递给 vectorize 装饰器。在较新版本的 numba 中，如果您将 numpy 数组作为输入传递给通用矢量化函数，则无需指定矢量输入类型。

例如：

import numpy as np
import numba as nb

a=np.array([1,2])
b=np.array([3,4])

# Note a generic vectorize decorator with input types not specified
@nb.vectorize
def f(x_1,x_2):
    return x_1+x_2

print(f(a,b))
#-> [4,6]

到目前为止，变量是从输入数组传递到函数的简单单个对象。这使得 numba 可以将 python 代码转换为可以对 numpy 数组进行操作的简单 ufunc。

在对向量求和的示例中，您需要将数据作为向量的单个向量传递。为此，您需要创建对向量本身进行操作的 ufunc。这需要更多的工作和规范，以了解如何创建任意输出。输入 guvectorize 函数 (docs here and here)。

因为您提供的是向量的向量。您的外部向量的处理方式类似于您在上面使用 vectorize 的方式。现在，您需要为输入值指定每个内部向量的外观。

EG 添加任意整数向量。（这将不起作用，原因如下所述）

@nb.guvectorize([(nb.int64[:])])
def f(x):
    return x.sum()

现在您还需要向函数和装饰器添加额外的输入。这允许您指定任意类型来存储函数的输出。您现在将更新此输入变量，而不是返回输出。将此最终变量视为 numba 在为 numpy 评估创建 ufunc 时用于生成任意输出向量的自定义变量。

这个输入也需要在装饰器中指定，你的函数应该是这样的：

@nb.guvectorize([(nb.int64[:],nb.int64[:])])
def f(x, out):
    out[:]=x.sum()

最后你需要在装饰器中指定输入和输出格式。这些以输入向量的顺序作为矩阵形状给出，并使用箭头指示输出向量形状（实际上是您的最终输入）。在这种情况下，您使用大小为 n 的向量并将结果输出为值而不是向量。您的格式应为 (n)->()。

作为一个更复杂的例子，假设您有两个大小为 (m,n) 和 (n,o) 的矩阵乘法的输入向量，并且您希望输出向量的大小为 (m,o) 您的装饰器格式看起来像(m,n),(n,o)->(m,o)。

当前问题的完整函数如下所示：

@nb.guvectorize([(nb.int64[:],nb.int64[:])], '(n)->()')
def f(x, out):
    out[:]=x.sum()

您的最终代码应该类似于：

import numpy as np
import numba as nb

a=np.arange(10).reshape(5,2)
# Equivalent to
# a=np.array([
#   [0,1],
#   [2,3],
#   [4,5],
#   [6,7],
#   [8,9]
# ])

@nb.guvectorize([(nb.int64[:],nb.int64[:])], '(n)->()')
def f(x, out):
    out[:]=x.sum()

print(f(a))
#-> [ 1  5  9 13 17]

当函数将数组作为参数时，如何使用Numba对函数进行矢量化处理？

2 个答案: