设置
我有以下两种矩阵计算的实现:
- 第一个实现使用形状为 (n, m) 的矩阵,并在 for 循环中重复计算
重复
次:
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
包含到矩阵中,该矩阵的形状为 (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 在处理三维而不是二维时速度较慢?
最佳答案
我不确定您的设置是什么,但我稍微重写了您的示例:
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 中的行为与您可能期望的不同:
通常最好像我在示例中那样传递数据。
第二种情况的问题是 numba 目前不支持 amin
。请参阅:
https://numba.pydata.org/numba-doc/dev/reference/numpysupported.html
如果将 nopython=True
传递给 jit
,您就可以看到这一点。因此,在当前版本的 numba(0.44 或更早版本)中,它将回退到 objectmode
,这通常并不比不使用 numba 快,有时还会更慢,因为存在一些调用开销。
关于python - 嵌套 Numpy 数组上的 Numba,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/56992398/