我知道numba会产生一些开销,并且在某些情况下(非密集计算),它变得比纯python慢。但是我不知道该划清界限。是否可以使用算法复杂度的顺序来找出位置?
例如,在此代码中,将两个数组(〜O(n))短于5的数组(〜O(n))更快:
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)
UPDADE:我正在寻找的是类似here的指南:
“一般准则是针对不同的数据大小和算法选择不同的目标。“ cpu”目标适用于较小的数据大小(大约小于1KB)和低计算强度的算法。其开销最少。“并行”目标适用于中等数据大小(约小于1MB);线程添加了较小的延迟;“ CUDA”目标适用于大数据大小(约大于1MB)和高计算强度算法。往返于GPU的内存转移会增加大量开销。”
答案 0 :(得分:3)
当numba生效时,很难划清界限。但是,有一些指标可能无效:
如果您不能将jit
与nopython=True
一起使用-每当您无法在nopython模式下进行编译时,您要么尝试编译太多,要么不会明显加快编译速度。
如果您不使用数组-当您处理传递给numba函数的列表或其他类型(其他numba函数除外)时,numba需要复制它们,这会产生大量开销。
如果已经有NumPy或SciPy函数可以执行此操作-即使numba对于短数组可以显着更快,但对于较长的数组几乎总是一样快(同样,您可能会轻易忽略一些常见的边缘情况,这些会处理)。
还有另一个原因,在某些情况下,您可能不希望使用numba只是比其他解决方案快一点:在某些情况下,必须提前编译Numba函数或在首次调用时进行编译即使您调用了数百次,编译也会比您获得的结果慢得多。编译时间也加起来: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解决方案会更好。
答案 1 :(得分:2)
如果您不完全知道显式输入和输出声明的结果,请让numba决定。输入您可能要使用'float64(float64[::1],float64[::1])'
。 (标量输出,连续的输入数组)。如果您使用跨步输入调用显式声明的函数,则该函数将失败,如果您使用Numba进行该工作,它将简单地重新编译。
如果不使用fastmath=True
,也将无法使用SIMD,因为它会改变结果的精度。
最好计算至少4个部分和(256位向量),而不是计算这些部分和的和(Numpy也不要计算幼稚的和)。
使用MSeiferts出色的基准测试实用程序的示例
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_zip(a,b):
result = 0.0
for i,j in zip(a,b):
result += (i+j)
return result
#Your version with suboptimal input and output (prevent njit compilation) declaration
@nb.jit('float64[:](float64[:],float64[:])')
def numba_your_func(a,b):
result = 0.0
for i,j in zip(a,b):
result += (i+j)
return result
@nb.njit(fastmath=True)
def numba_loop_zip_fastmath(a,b):
result = 0.0
for i,j in zip(a,b):
result += (i+j)
return result
@nb.njit(fastmath=True)
def numba_loop_fastmath_single(a,b):
result = 0.0
size=min(a.shape[0],b.shape[0])
for i in range(size):
result += a[i]+b[i]
return result
@nb.njit(fastmath=True,parallel=True)
def numba_loop_fastmath_multi(a,b):
result = 0.0
size=min(a.shape[0],b.shape[0])
for i in nb.prange(size):
result += a[i]+b[i]
return result
#just for fun... single-threaded for small arrays,
#multithreaded for larger arrays
@nb.njit(fastmath=True,parallel=True)
def numba_loop_fastmath_combined(a,b):
result = 0.0
size=min(a.shape[0],b.shape[0])
if size>2*10**4:
result=numba_loop_fastmath_multi(a,b)
else:
result=numba_loop_fastmath_single(a,b)
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, 19)
}
b = benchmark([python_loop, numba_loop_zip, numpy_methods,numba_your_func, python_sum,numba_loop_zip_fastmath,numba_loop_fastmath_single,numba_loop_fastmath_multi,numba_loop_fastmath_combined], arguments, warmups=[numba_loop_zip,numba_loop_zip_fastmath,numba_your_func,numba_loop_fastmath_single,numba_loop_fastmath_multi,numba_loop_fastmath_combined])
%matplotlib notebook
b.plot()
请注意,仅在某些特殊情况下建议使用numba_loop_fastmath_multi
或numba_loop_fastmath_combined(a,b)
。这种简单的功能通常是另一个问题的一部分,可以更有效地并行化(启动线程有一些开销)
答案 2 :(得分:1)
运行此代码会导致我的计算机上的速度提高约6倍:
@numba.autojit
def sum_2(a,b):
result = 0.0
for i,j in zip(a,b):
result += (i+j)
return result
Python:3.31 µs,数字:589 ns。
关于您的问题,我真的认为这与复杂性并没有真正的关系,它可能主要取决于您正在执行的操作类型。另一方面,您仍然可以绘制python / numba比较图,以查看给定函数发生移位的地方。