Numba自动jit函数比向量化的Numpy方法慢

Question

我有以下的for循环来创建一个值列表：

p = 7
A = []

for i in range(0, 10**p):
    A.append(i**3 + i**2)

为了加快列表的创建速度，我使用向量化方法将其创建为Numpy数组。这种方法比等效的for循环更快，特别是对于增加范围的大值p。

import numpy as np
from numba import autojit

p = 7
m = np.arange(0, 10**p)
D = np.empty(len(m))
D = m**3 + m**2

为了进一步加快创建数组的速度，我考虑尝试使用Numba软件包。以下是我的当前尝试。

@autojit
def func(a):
    a = np.asarray(a)
    n = np.arange(0, 10**p)
    a = np.append(a, n**3 + n**2)
    return a

e = []
E = func(e)

很不幸，使用Numba并没有带来性能提升，而且速度几乎比仅使用Numpy的向量化方法慢3倍。

有没有关于如何使用Numba的建议？

- wigging

---

numpy函数已经向量化，使用@autojit将在处理for循环时产生巨大的差异。 - Saullo G. P. Castro

@SaulloCastro，你的意思是将for循环封装到一个带有autojit的函数中？你能提交一个使用这种方法的答案并比较时间吗？ - wigging

我的答案与JoshAdels的类似，重点是在使用Numba时像Fortran一样编写代码，只是提醒这里的索引从0开始 :) - Saullo G. P. Castro

Answer 1

Numba无法使任意方法调用更快。如果您调用库，大多数情况下numba真的无能为力。但是，如果您稍微改写一下代码，仍然可以获得不错的加速效果（我正在使用numba 0.14.0——如果您使用的是不同版本、硬件等，可能会得到不同的结果，特别是因为numba正在积极开发中）：

import numpy as np
import numba as nb

def func(a, p):
    a = np.asarray(a)
    n = np.arange(0, 10**p)
    a = np.append(a, n**3 + n**2)
    return a

@nb.jit
def func2(a, p):
    a = np.asarray(a)
    n = np.empty(10**p, dtype=np.float64)
    for k in range(10**p):
        n[k] = k*k*(k + 1)

    return np.append(a, n)

p = 6
e = []
E = func(e, p)
E2 = func2(e, p)
print np.allclose(E, E2)

并且包括时间安排：

In [51]:

%timeit func(e, p)
10 loops, best of 3: 42.9 ms per loop
In [52]:

%timeit func2(e, p)
100 loops, best of 3: 3.09 ms per loop

如果使用p=7，需要注意数值精度问题。

Numba的关键是展开循环，并只进行numba支持的“原始”算术调用。

- JoshAdel

---

你的func2中速度的提升是由于k*k*(k+1)而不是仅仅使用nb.jit。如果你在func中将n**3+n**2改为n*n*(n+1)，那么它仍然会比使用numba方法更快。 - wigging

在我的计算机上， n*n*(n+1) 的速度仍然比numba慢2倍。实际上我有点惊讶，因为我认为创建的临时数组数量应该是相似的，但我猜测 ** 运算符的代码路径可能有些不同。 - JoshAdel

也许我的代码还有其他问题。这是我正在运行的文件 https://github.com/wigging/python_examples/blob/master/vectorize.py - wigging

@Gavin 我猜你只计时了 Numba 函数一次，所以时间包括了代码的 JIT 编译时间。而当我使用 timeit 时，我从一个已经预热过的状态获得最佳时间。 - JoshAdel

最新版本似乎为@jit引入了一个“cache”标志，这是版本0.22。当numba不再是一个移动目标时，这将非常令人兴奋... - webb