嵌套 Numpy 数组上的 Numba

Question

设置

我有以下两个矩阵计算的实现：

第一个实现使用 matrix of shape (n, m) 并在 for 循环中重复计算 repetition 次：

import numpy as np
from numba import jit

@jit
def foo():
    for i in range(1, n):
        for j in range(1, m):

            _deleteA = (
                        matrix[i, j] +
                        #some constants added here
            )
            _deleteB = (
                        matrix[i, j-1] +
                        #some constants added here
            )
            matrix[i, j] = min(_deleteA, _deleteB)

    return matrix

repetition = 3
for x in range(repetition):
    foo()

2. 第二种实现避免了额外的 for 循环，因此将 repetition = 3 包含到矩阵中，然后是 shape (repetition, n, m):

@jit
def foo():
    for i in range(1, n):
        for j in range(1, m):

            _deleteA = (
                        matrix[:, i, j] +
                        #some constants added here
            )
            _deleteB = (
                        matrix[:, i, j-1] +
                        #some constants added here
            )
            matrix[:, i, j] = np.amin(np.stack((_deleteA, _deleteB), axis=1), axis=1)

    return matrix

问题

关于这两个实现，我在 iPython.

中发现了两点关于它们在 %timeit 上的表现

第一个实现从 @jit 中获利巨大，而第二个则根本没有（在我的测试用例中为 28 毫秒对 25 秒）。 谁能想象为什么 @jit 不再适用于形状为 (repetition, n, m) 的 numpy 数组？

编辑

我将之前的第二个问题移到了 an extra post，因为问多个问题被认为是糟糕的 SO 风格。

问题是：

当忽略 @jit 时，第一个实现仍然快很多（相同的测试用例：17 秒对 26 秒）。 为什么 numpy 在处理三维而不是二维时速度较慢？

Answer 1

我不确定你这里的设置是什么，但我re-wrote你的例子略：

import numpy as np
from numba import jit

#@jit(nopython=True)
def foo(matrix):
    n, m = matrix.shape
    for i in range(1, n):
        for j in range(1, m):

            _deleteA = (
                        matrix[i, j] #+
                        #some constants added here
            )
            _deleteB = (
                        matrix[i, j-1] #+
                        #some constants added here
            )
            matrix[i, j] = min(_deleteA, _deleteB)

    return matrix

foo_jit = jit(nopython=True)(foo)

然后时间：

m = np.random.normal(size=(100,50))

%timeit foo(m)  # in a jupyter notebook
# 2.84 ms ± 54.2 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%timeit foo_jit(m)  # in a jupyter notebook
# 3.18 µs ± 38.9 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

所以这里的 numba 比预期的要快很多。需要考虑的一件事是全局 numpy 数组在 numba 中的行为并不像您预期的那样：

https://numba.pydata.org/numba-doc/dev/user/faq.html#numba-doesn-t-seem-to-care-when-i-modify-a-global-variable

通常像我在示例中那样传入数据会更好。

第二种情况下你的问题是numba目前不支持amin。参见：

https://numba.pydata.org/numba-doc/dev/reference/numpysupported.html

如果将 nopython=True 传递给 jit 就可以看到这个。因此，在当前版本的 numba（当前为 0.44 或更早版本）中，它会回退到 objectmode，这通常并不比不使用 numba 快，而且有时会更慢，因为存在一些调用开销。

嵌套 Numpy 数组上的 Numba

Numba on nested Numpy arrays

python

arrays

numpy

matrix

numba