【问题标题】:How do I vectorize a function with Numba, when the function takes arrays as arguments?当函数将数组作为参数时,如何使用 Numba 对函数进行矢量化?
【发布时间】:2018-12-24 11:08:09
【问题描述】:

我想使用 Numba 来向量化一个函数,该函数将评估矩阵的每一行。这实际上会将 Numpy ufunc 应用于矩阵,而不是循环遍历行。根据docs

您可能会问自己,“我为什么要通过这个而不是使用 @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 只接受标量输入。我还编写了签名以获取数组参数并修改标量输出。根据docs,装饰函数不使用return 语句。

结果应该是b = [1,5,9,13,17],但我得到的是b=[0.,1.,2.,3.,4.]。显然,我错过了一些东西。我会很感激一些方向,记住总和只是一个例子。

【问题讨论】:

    标签: python numpy scipy numba


    【解决方案1】:

    @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 hereand 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]
    

    【讨论】:

      【解决方案2】:

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

      【讨论】:

      • 非常有趣!谢谢!我不认为我会自己解决这个问题。
      猜你喜欢
      • 1970-01-01
      • 1970-01-01
      • 1970-01-01
      • 1970-01-01
      • 1970-01-01
      • 2020-01-11
      • 1970-01-01
      • 1970-01-01
      • 1970-01-01
      相关资源
      最近更新 更多