Numba 对具有类型参数的函数的无效使用答案

【问题标题】：Numba Invalid use of Function with argument(s) of type(s)Numba 对具有类型参数的函数的无效使用
【发布时间】：2020-05-19 07:31:34
【问题描述】：

我正在使用 Numba 非 Python 模式和一些 NumPy 函数。

@njit
def invert(W, copy=True):
    '''
    Inverts elementwise the weights in an input connection matrix.
    In other words, change the from the matrix of internode strengths to the
    matrix of internode distances.

    If copy is not set, this function will *modify W in place.*

    Parameters
    ----------
    W : np.ndarray
        weighted connectivity matrix
    copy : bool

    Returns
    -------
    W : np.ndarray
        inverted connectivity matrix
    '''

    if copy:
        W = W.copy()
    E = np.where(W)
    W[E] = 1. / W[E]
    return W

在这个函数中，W 是一个矩阵。但我收到以下错误。它可能与W[E] = 1. / W[E] 行有关。

File "/Users/xxx/anaconda3/lib/python3.7/site-packages/numba/dispatcher.py", line 317, in error_rewrite
    reraise(type(e), e, None)
  File "/Users/xxx/anaconda3/lib/python3.7/site-packages/numba/six.py", line 658, in reraise
    raise value.with_traceback(tb)
numba.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Invalid use of Function(<built-in function getitem>) with argument(s) of type(s): (array(float64, 2d, A), tuple(array(int64, 1d, C) x 2))

那么使用 NumPy 和 Numba 的正确方法是什么？我知道 NumPy 在矩阵计算方面做得很好。在这种情况下，NumPy 是否足够快以至于 Numba 不再提供加速？

【问题讨论】：

Numba 不支持“花式”索引，请查看 here。您应该沿着两个数组维度循环

标签： python numpy numba

【解决方案1】：

正如 FBruzzesi 在 cmets 中提到的，代码未编译的原因是您使用了“花式索引”，因为 W[E] 中的 E 是 np.where 的输出，并且是数组的元组。（这解释了稍微神秘的错误消息：Numba 不知道如何使用 getitem，即当输入之一是元组时，它不知道如何在括号中查找内容。）

Numba actually supports fancy indexing (also called "advanced indexing") on a single dimension，只是不是多个维度。在您的情况下，这允许进行简单的修改：首先使用ravel 几乎毫无成本地使您的数组一维，然后应用转换，然后使用廉价的reshape 返回。

@njit
def invert2(W, copy=True):
    if copy:
        W = W.copy()
    Z = W.ravel()
    E = np.where(Z)
    Z[E] = 1. / Z[E]
    return Z.reshape(W.shape)

但这仍然比它需要的要慢，因为通过不必要的中间数组传递计算，而不是在遇到非零值时立即修改数组。简单地做一个循环会更快：

@njit 
def invert3(W, copy=True): 
    if copy: 
        W = W.copy() 
    Z = W.ravel() 
    for i in range(len(Z)): 
        if Z[i] != 0: 
            Z[i] = 1/Z[i] 
    return Z.reshape(W.shape)

无论W 的尺寸如何，此代码都有效。如果我们知道W 是二维的，那么我们可以直接在这两个维度上进行迭代，但是由于两者具有相似的性能，我将采用更通用的路线。

在我的电脑上，假设一个 300×300 数组 W 大约一半的条目是 0，而 invert 是您在没有 Numba 编译的情况下的原始函数，时间安排是：

In [80]: %timeit invert(W)                                                                   
2.67 ms ± 49.3 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [81]: %timeit invert2(W)                                                                  
519 µs ± 24.5 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

In [82]: %timeit invert3(W)                                                                  
186 µs ± 11.1 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

所以 Numba 为我们提供了相当可观的加速（在它已经运行一次以消除编译时间之后），尤其是在以 Numba 可以利用的高效循环样式重写代码之后。

【讨论】：