我正在尝试使用 numba.jit 加速一个简单的 Python 循环。但似乎 jit 无法处理基本的数组索引和切片？我能做些什么来使这成为可能吗？如果 jit 不能处理基本的 numpy 数组，我不明白它怎么会有用。
我强制使用 nopython 模式。它在对象模式下工作，但这根本不会加速代码，所以它是我需要的 nopython 模式。
下面的代码只是一个说明问题的例子。我的实际代码有同样的问题，但有更多的循环和迭代，因此 jit 会非常有用。
import numpy as np
from numba import jit
n = 100
myarray = np.zeros(n)
@jit(nopython=True)
def compute(n):
for i in xrange(n):
myarray[i] += 1 # This indexing causes the error.
compute(n)
"""Sample run:
> python jit_test.py
> ...
> Failed at nopython (nopython frontend)
> Internal error at <numba.typeinfer.SetItemConstrain object at 0x7f700c89a7d0>:
> Immutable array
> File "jit_test.py", line 10
"""
如果我尝试切片而不是出现不同的错误。
# The rest of the code unchanged.
myarray[:] += 1 # This slicing causes the error.
# The rest of the code unchanged.
"""Sample run:
> python jit_test.py
> ...
> Failed at nopython (nopython frontend)
> Internal error at <numba.typeinfer.IntrinsicCallConstrain object at 0x7f297e2a9dd0>:
> Don't know how to create implicit output array with 'A' layout.
> File "jit_test2.py", line 10
"""
我的构建:
numba version: 0.19.1
numpy version: 1.9.2
python version: 2.7.10
请您参考如下方法:
在 numba 中，全局数组是静态的。您将在您的方法中获得一个只读副本。见:
如果您想修改一个数组，要么显式传入它，要么在函数中创建它(在进行数组内存管理的较新版本的 Numba 中)并返回它。
例如:
import numpy as np
import numba as nb
def compute(arr):
for i in xrange(arr.shape[0]):
arr[i] += 1
n = 100
myarray = np.zeros(n)
jitcompute = nb.jit(nopython=True)(compute)
然后计时:
In [12]: %timeit compute(myarray)
10000 loops, best of 3: 25.7 µs per loop
In [13]: %timeit jitcompute(myarray)
The slowest run took 17.06 times longer than the fastest. This could mean that an intermediate result is being cached
1000000 loops, best of 3: 461 ns per loop
一些注意事项:
- 这显然是一个简单的示例，但它展示了如何在方法之外定义一个数组，然后让函数就地修改它。
- 我没有在这里使用
@jit装饰器，这样我就可以编写一次方法定义，并在使用和不使用 numba jit-ing 的情况下为方法计时，但是jitcompute等同于在compute上使用装饰器会得到的结果。 - 如果您曾经对 numba 代码进行基准测试，请确保在计时之前运行该方法一次，否则您将看到的是 jit 代码所花费的时间和实际执行时间的组合。第二次运行时，您只会看到执行时间。
