Numba的autojit函数比向量化的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的建议？

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支持的“基本”算术运算。