Numba：为什么 guvectorize 这么慢？答案

【问题标题】：Numba : Why guvectorize is so slow?Numba：为什么 guvectorize 这么慢？
【发布时间】：2022-01-17 01:58:24
【问题描述】：

我正在学习 numba 以尝试优化我的代码。首先，我计算了 a*exp(b*x) 其中 a 和 b 是两个参数，x 是一个大的 numpy 数组。我使用了@jit 装饰器。输出是具有 len(x) 个元素的一维数组。这是代码：

import timeit

import numpy as np
from numba import jit, prange

def func_np (par,x):
    return(par[0]*np.exp(par[1]*x))

@jit('float64[:](float64[:], float64[:])', nopython=True, parallel=True)
def func_parallel(par,x):
    length=len(x)
    result = np.empty(length, dtype=np.float64)
    for i in prange(length):
        result[i] = par[0]*np.exp(par[1]*x[i])
    return result

x=np.array(np.arange(0,100,0.0001))
par=np.array([10,0.1])
print("numpy only")
%timeit func_np(par,x)
print("")
print("numba")
%timeit func_parallel(par,x)

输出是：

numpy only
28.1 ms ± 27.5 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

numba
2.05 ms ± 13.4 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

所以，一切都很顺利。通过 numba 优化，我的时间减少了很多。（运行此代码的机器有 48 个内核。）。

之后，我决定使用 a 和 b 参数的数组（par=np.array([[i*10,i*0.1] for i in range(10)]），x 不变。这是我的问题的开始。首先，由于输出是形状 (10 ,len(x)) 元素的二维数组，我似乎不能使用 @jit 装饰器，但我必须使用 @guvectorize 装饰器。我编写了这段代码（与上面的代码非常相似）：

import timeit

import numpy as np
from numba import guvectorize, prange

def func_np (par,x):
    return(par[:,0:1]*np.exp(par[:,1:2]*x))

@guvectorize([ 'float64[:,:], float64[:,:], float64[:]' ], '(m,n),(m,p),(n)', nopython=True, target='parallel')
def func_parallel(result,par,x):
    lenght=len(x)
    for i in prange(lenght):
        result[:,i] = (par[:,0:1]*np.exp(par[:,1:2]*x[i])).T
         
x=np.array(np.arange(0,100,0.0001))
par=np.array([[i*10,i*0.1] for i in range(10)])
print("numpy only")
%timeit func_np(par,x)
print("")
print("numba")
result = np.empty([par.shape[0],x.shape[0]], dtype=np.float64)
%timeit func_parallel(result, par,x);

但是，基准并不好。使用 numba 优化的代码比仅使用 numpy 函数编写的代码要慢。

numpy only
358 ms ± 307 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

numba
563 ms ± 2.82 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

因为我是第一次使用 numba，所以我想有一些我没有理解的东西。任何帮助将不胜感激。

【问题讨论】：

标签： python arrays numpy numba

【解决方案1】：

我不完全知道 guvectorize 会做什么来减慢您的速度（很可能会创建额外的不必要的线程），但很容易做出更好定义的实现。

考虑以下比 numpy 实现更快的实现。

@jit('float64[:,:](float64[:,:], float64[:])', nopython=True,parallel=True)
def func_parallel(par, x):
    result = np.zeros((par.shape[0],len(x)),dtype=np.float64)
    lenght = len(x)
    for i in prange(lenght):
        for j in range(par.shape[0]):
            result[j, i] = par[j, 0] * np.exp(par[j, 1] * x[i])
    return result

在我的 2C/4T 机器上运行速度要快 3 倍，所以在你的机器上应该快得多。

【讨论】：