Python中Numba jit警告的解释

11
我定义了以下递归数组生成器,并使用Numba jit尝试加速处理(基于此SO答案)。
@jit("float32[:](float32,float32,intp)", nopython=False, nogil=True)
def calc_func(a, b, n):
    res = np.empty(n, dtype="float32")
    res[0] = 0
    for i in range(1, n):
        res[i] = a * res[i - 1] + (1 - a) * (b ** (i - 1))
    return res
a = calc_func(0.988, 0.9988, 5000)

我收到了一堆警告/错误,但并不太明白。希望能得到帮助来解释它们,并使它们消失,以加快计算速度。

以下是这些警告/错误:

NumbaWarning: 编译正在回退到启用循环提升的对象模式,因为“calc_func”函数由于以下原因未通过类型推断:使用无效的Function()参数类型:(int64, dtype=Literalstr) * 参数化

在定义0中: 所有模板都被字面量拒绝。

在定义1中: 所有模板都没有字面量被拒绝。 这个错误通常是由于传递了一个不受命名函数支持的类型的参数而引起的。

[1] 在:解析调用者类型时

[2] 在:res = np.empty(n, dtype="float32")的调用类型

文件“thenameofmyscript.py”,第71行:

def calc_func(a, b, n):
    res = np.empty(n, dtype="float32")
    ^

@jit("float32:", nopython=False, nogil=True)

thenameofmyscript.py:69: NumbaWarning: 编译回退到对象模式,因为函数“calc_func”的类型推断失败,原因是无法确定<class 'numba.dispatcher.LiftedLoop'>的Numba类型

File "thenameofmyscript.py", line 73:

def calc_func(a, b, n):
        <source elided>
        res[0] = 0
        for i in range(1, n):
        ^

@jit("float32:", nopython=False, nogil=True)

H:\projects\decay-optimizer\venv\lib\site-packages\numba\compiler.py:742: NumbaWarning: 函数 "calc_func" 在未使用 forceobj=True 的情况下以对象模式编译,但已升级循环。

文件 "thenameofmyscript.py",行 70:

@jit("float32[:](float32,float32,intp)", nopython=False, nogil=True)
    def calc_func(a, b, n):
    ^

self.func_ir.loc))

H:\projects\decay-optimizer\venv\lib\site-packages\numba\compiler.py:751: NumbaDeprecationWarning: 检测到从nopython编译路径回退到对象模式编译路径,这是不推荐的行为。

文件"thenameofmyscript.py",第70行:

@jit("float32[:](float32,float32,intp)", nopython=False, nogil=True)
    def calc_func(a, b, n):
    ^

警告:warnings.warn(errors.NumbaDeprecationWarning(msg, self.func_ir.loc))

thenameofmyscript.py:69: NumbaWarning: 代码在对象模式下运行,即使nogil=True也不允许并行执行。 @jit("float32:", nopython=False, nogil=True)


1
  1. 使用njit或nopython=True -> 这将导致错误(回退到对象模式已被弃用)。
  2. 这是错误的语法:np.empty(n, dtype="float32") 将其更改为np.empty(n,dtype=np.float32),正如您在numpy中通常所做的那样。
  3. 通常不需要指定数据类型。 您可以完全省略此声明。 "float32[:](float32,float32,intp)"
- max9111
做到了,非常感谢。给函数签名会让它更快吗?类型推断已经足够好了吗?如果你想发表答案,我会验证它。 - Chapo
不需要签名。这个问题与此 https://stackoverflow.com/a/57062221/4045774 有些相关(昂贵的指数可能是可避免的),但我强烈建议使用float64。 - max9111
1个回答

6

1. 优化函数(代数简化)

现代CPU对于加法、减法和乘法的处理速度非常快。尽可能避免使用指数运算等操作。

例如:

在本例中,我把耗时的指数运算替换成了简单的乘法。像这样的简化可以大幅提高运行速度,但也可能改变结果。

首先,您的实现(float64)没有任何标识,稍后我会在另一个简单的示例中进行处理。

#nb.jit(nopython=True) is a shortcut for @nb.njit()
@nb.njit()
def calc_func_opt_1(a, b, n):
    res = np.empty(n, dtype=np.float64)
    fact=b
    res[0] = 0.
    res[1] = a * res[0] + (1. - a) *1.
    res[2] = a * res[1] + (1. - a) * fact
    for i in range(3, n):
        fact*=b
        res[i] = a * res[i - 1] + (1. - a) * fact
    return res

尽可能使用标量变量是一个好主意。

@nb.njit()
def calc_func_opt_2(a, b, n):
    res = np.empty(n, dtype=np.float64)
    fact_1=b
    fact_2=0.
    res[0] = fact_2
    fact_2=a * fact_2 + (1. - a) *1.
    res[1] = fact_2
    fact_2 = a * fact_2 + (1. - a) * fact_1
    res[2]=fact_2
    for i in range(3, n):
        fact_1*=b
        fact_2= a * fact_2 + (1. - a) * fact_1
        res[i] = fact_2
    return res

时间表

%timeit a = calc_func(0.988, 0.9988, 5000)
222 µs ± 2.2 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
%timeit a = calc_func_opt_1(0.988, 0.9988, 5000)
22.7 µs ± 45.5 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
%timeit a = calc_func_opt_2(0.988, 0.9988, 5000)
15.3 µs ± 35.6 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

2. 签名是否推荐?

在提前编译模式(AOT)下,签名是必需的,但在通常的JIT模式下不需要。上面的示例不支持SIMD向量化。因此,您可能不会看到关于输入和输出可能不太优化的声明的太多积极或消极影响。让我们看另一个例子。

#Numba is able to SIMD-vectorize this loop if 
#a,b,res are contigous arrays
@nb.njit(fastmath=True)
def some_function_1(a,b):
    res=np.empty_like(a)
    for i in range(a.shape[0]):
        res[i]=a[i]**2+b[i]**2
    return res

@nb.njit("float64[:](float64[:],float64[:])",fastmath=True)
def some_function_2(a,b):
    res=np.empty_like(a)
    for i in range(a.shape[0]):
        res[i]=a[i]**2+b[i]**2
    return res

a=np.random.rand(10_000)
b=np.random.rand(10_000)

#Example for non contiguous input
#a=np.random.rand(10_000)[0::2]
#b=np.random.rand(10_000)[0::2]

%timeit res=some_function_1(a,b)
5.59 µs ± 36.1 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
%timeit res=some_function_2(a,b)
9.36 µs ± 47.1 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

为什么带签名的版本会更慢?
让我们仔细看一下签名。
some_function_1.nopython_signatures
#[(array(float64, 1d, C), array(float64, 1d, C)) -> array(float64, 1d, C)]
some_function_2.nopython_signatures
#[(array(float64, 1d, A), array(float64, 1d, A)) -> array(float64, 1d, A)]
#this is equivivalent to 
#"float64[::1](float64[::1],float64[::1])"

如果内存布局在编译时未知,通常无法对算法进行SIMD向量化。当然,您可以显式声明C连续数组,但是该函数将不再适用于非连续输入,这通常不是预期的。


非常感谢。我不明白为什么使用标量会有所不同。从视觉上看,它看起来像是在进行更多的操作,但结果却完全相同。它保留了fact_2在内存中的事实造成了差异吗? - Chapo

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接