Question

我是numba的jit的新手。对于个人项目，我需要加速与下面显示的功能类似的功能，但为了编写独立示例的目的不同。

import numpy as np
from numba import jit, autojit, double, float64, float32, void

def f(n):
    k=0.
    for i in range(n):
        for j in range(n):
            k+= i+j

def f_with_return(n):
    k=0.
    for i in range(n):
        for j in range(n):
            k+= i+j
    return k

def f_with_arange(n):
    k=0.
    for i in np.arange(n):
        for j in np.arange(n):
            k+= i+j

def f_with_arange_and_return(n):
    k=0.
    for i in np.arange(n):
        for j in np.arange(n):
            k+= i+j  


#jit decorators
jit_f = jit(void(int32))(f)
jit_f_with_return = jit(int32(int32))(f_with_return)
jit_f_with_arange = jit(void(double))(f_with_arange)
jit_f_with_arange_and_return = jit(double(double))(f_with_arange_and_return)

基准：

%timeit f(1000)
%timeit jit_f(1000)

10个循环，最佳3：每循环73.9 ms / 1000000循环，最佳3：212 ns /循环

%timeit f_with_return(1000)
%timeit jit_f_with_return(1000)

10个循环，最佳3：每循环74.9 ms / 1000000循环，最佳3：每循环220 ns

我不明白这两个：

%timeit f_with_arange(1000.0)
%timeit jit_f_with_arange(1000.0)

10个循环，最佳3：175 ms /循环/ 1循环，最佳3：167 ms每循环

%timeit f_with_arange_with_return(1000.0)
%timeit jit_f_with_arange_with_return(1000.0)

10个循环，最佳3：174 ms每循环/ 1循环，最佳3：172 ms每循环

我认为我没有给jit函数输出和输入的正确类型？只是因为for循环现在在numpy.arange上运行，而不再是一个简单的范围，我不能让jit让它更快。这是什么问题？

Answer 1

简单地说，numba不知道如何将np.arange转换为低级本机循环，因此它默认为对象层，它比较慢，通常与纯python的速度相同。

一个很好的技巧是将nopython=True关键字参数传递给jit以查看它是否可以编译所有内容而无需使用对象模式：

import numpy as np
import numba as nb

def f_with_return(n):
    k=0.
    for i in range(n):
        for j in range(n):
            k+= i+j
    return k

jit_f_with_return = nb.jit()(f_with_return)
jit_f_with_return_nopython = nb.jit(nopython=True)(f_with_return)

%timeit f_with_return(1000)
%timeit jit_f_with_return(1000)
%timeit jit_f_with_return_nopython(1000)

最后两个在我的机器上的速度相同，并且比未校对的代码快得多。您遇到问题的两个示例会引发nopython=True错误，因为此时无法编译np.arange。

有关详细信息，请参阅以下内容：

http://numba.pydata.org/numba-doc/0.17.0/user/troubleshoot.html#the-compiled-code-is-too-slow

以及支持的numpy功能列表，其中包含nopython模式支持和不支持的内容的指示：

http://numba.pydata.org/numba-doc/0.17.0/reference/numpysupported.html

使用numba JIT加速功能的麻烦

1 个答案: