numba什么时候有效？

Question

我知道 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:

"A general guideline is to choose different targets for different data sizes and algorithms. The “cpu” target works well for small data sizes (approx. less than 1KB) and low compute intensity algorithms. It has the least amount of overhead. The “parallel” target works well for medium data sizes (approx. less than 1MB). Threading adds a small delay. The “cuda” target works well for big data sizes (approx. greater than 1MB) and high compute intensity algorithms. Transfering memory to and from the GPU adds significant overhead."

Answer 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 微秒，数字：589 纳秒。

至于你的问题，我真的认为这与复杂性并没有真正的关系，它可能主要取决于你正在进行的操作类型。另一方面，您仍然可以绘制一个 python/numba 比较，以查看给定函数发生偏移的位置。

Answer 2

很难说 numba 何时生效。但是，有一些指标可能无效：

如果您不能将 jit 与 nopython=True 一起使用 - 每当您无法在 nopython 模式下编译它时，您要么尝试编译太多，要么它不会明显更快。
如果您不使用数组 - 当您处理传递给 numba 函数的列表或其他类型时（其他 numba 函数除外），numba 需要复制这些，这会导致显着开销。
如果已经有一个 NumPy 或 SciPy 函数可以做到这一点 - 即使 numba 对于短数组来说可以快得多，但对于更长的数组来说几乎总是一样快（你也是可能很容易忽略这些可以处理的一些常见边缘情况。

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

而且 numba 安装起来很复杂（至少没有 conda）所以如果你想分享你的代码那么你有一个真正的 "heavy dependency".

您的示例缺少与 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 解决方案较慢，但 "faster" 替代方案已经比您最初提出的解决方案快得多。

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

Answer 3

如果您不确切知道显式输入和输出声明的结果是什么，请让 numba 决定。根据您的输入，您可能希望使用 'float64(float64[::1],float64[::1])'。（标量输出，连续输入数组）。如果您使用跨步输入调用显式声明的函数，它将失败，如果您让 Numba 完成这项工作，它只会重新编译。不使用 fastmath=True 也不可能使用 SIMD，因为它会改变结果的精度。

计算至少 4 个部分和（256 位向量）并且比计算这些部分和的总和在这里更可取（Numpy 也不计算朴素的总和）。

使用 MSeifert 优秀基准实用程序的示例

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)。通常这样一个简单的函数是另一个问题的一部分，可以更有效地并行化（启动线程有一些开销）

numba什么时候有效？

When numba is effective?

python

performance

python-3.x

numba