Skip to main content
 首页 » 编程设计

python之numba什么时候有效

2025年05月04日15落叶

我知道 numba 会产生一些开销,并且在某些情况下(非密集计算)它变得比纯 python 慢。但我不知道在哪里划清界限。是否可以使用算法复杂度的顺序来找出位置?

例如,在这段代码中添加两个短于 5 的数组 (~O(n)) 纯 python 更快:

def sum_1(a,b): 
    result = 0.0 
    for i,j in zip(a,b): 
            result += (i+j) 
    return result 
 
@numba.jit('float64[:](float64[:],float64[:])') 
def sum_2(a,b): 
    result = 0.0 
    for i,j in zip(a,b): 
            result += (i+j) 
    return result 
 
# try 100 
a = np.linspace(1.0,2.0,5) 
b = np.linspace(1.0,2.0,5) 
print("pure python: ") 
%timeit -o sum_1(a,b) 
print("\n\n\n\npython + numba: ") 
%timeit -o sum_2(a,b) 

更新:我正在寻找类似的指南,如 here :

“一般准则是为不同的数据大小和算法选择不同的目标。“cpu”目标适用于小数据大小(大约小于 1KB)和低计算强度算法。它的开销最少. “并行”目标适用于中等数据大小(大约小于 1MB)。线程增加了一个小的延迟。“cuda”目标适用于大数据大小(大约大于 1MB)和高计算强度算法。将内存传入和传出 GPU 会增加大量开销。”

请您参考如下方法:

当 numba 生效时很难划清界线。但是,有一些指标可能不有效:

  • 如果您不能将 jit 与 nopython=True 一起使用 - 每当您无法在 nopython 模式下编译它时,您要么尝试编译太多,要么它不会明显更快。

  • 如果您不使用数组 - 当您处理传递给 numba 函数的列表或其他类型时(其他 numba 函数除外),numba 需要复制这些会产生大量开销。

  • 如果已经有一个 NumPy 或 SciPy 函数可以执行此操作 - 即使 numba 对于短数组可以明显更快,但对于更长的数组几乎总是一样快(而且您可能很容易忽略一些常见的边缘情况这些会处理)。

在 numba 比其他解决方案“快一点”的情况下,您可能不想使用 numba 还有另一个原因:Numba 函数必须提前编译或编译第一次调用时,在某些情况下编译会比你的 yield 慢得多,即使你调用它数百次。编译时间也会增加:numba 导入速度很慢,编译 numba 函数也会增加一些开销。如果导入开销增加 1-10 秒,则削减几毫秒是没有意义的。

而且 numba 安装起来很复杂(至少没有 conda),所以如果你想分享你的代码,那么你就会有一个真正的“严重依赖”。


您的示例缺少与 NumPy 方法和高度优化的纯 Python 版本的比较。我添加了更多比较函数并进行了基准测试(使用我的库 simple_benchmark ):

import numpy as np 
import numba as nb 
from itertools import chain 
 
def python_loop(a,b): 
    result = 0.0 
    for i,j in zip(a,b): 
        result += (i+j) 
    return result 
 
@nb.njit 
def numba_loop(a,b): 
    result = 0.0 
    for i,j in zip(a,b): 
            result += (i+j) 
    return result 
 
def numpy_methods(a, b): 
    return a.sum() + b.sum() 
 
def python_sum(a, b): 
    return sum(chain(a.tolist(), b.tolist())) 
 
from simple_benchmark import benchmark, MultiArgument 
 
arguments = { 
    2**i: MultiArgument([np.zeros(2**i), np.zeros(2**i)]) 
    for i in range(2, 17) 
} 
b = benchmark([python_loop, numba_loop, numpy_methods, python_sum], arguments, warmups=[numba_loop]) 
 
%matplotlib notebook 
b.plot() 

是的,对于小数组,numba 函数是最快的,但是对于较长的数组,NumPy 解决方案会稍微快一些。 Python 解决方案速度较慢,但​​“更快”的替代方案已经比您最初提出的解决方案快得多。

在这种情况下,我会简单地使用 NumPy 解决方案,因为它简短、可读且速度快,除非您处理大量短数组并多次调用该函数 - 那么 numba 解决方案会好得多。