numpy Numba -有没有可能做绝对和没有nan更快?

2ul0zpep  于 2023-03-02  发布在  其他
关注(0)|答案(1)|浏览(113)

我有一个稍微修改的例子从numba官方文件如下:

from numba import njit
import numpy as np

@njit
def do_sum(A, lb, ub):
    n = len(A)
    acc = 0.0
    for i in range(n):
        a = 0.0 if np.isnan(A[i]) else A[i]
        acc += abs(max(min(a, ub[i]), lb[i]))
    return acc

基本上,任务是对一个numpy数组进行绝对有界求和,这个数组可能有np.nanAlbub都是长度相同的1d数组。长度通常在数千或数万的范围内。而且,函数预计会被调用数百万次-这就是为什么我想优化它。
由于A可能包含np.nan,所以我似乎不能使用@njit(fastmath=True),但是,当A不包含np.nan时,我的基准测试结果显示,使用@njit(fastmath=True)明显比@njit快。
我的问题是,在这两者之间是否存在一些最佳点,这样我就可以使代码与fastmath一起工作,并获得比上面的do_sum实现更快的速度?
或者说,无论什么样的方法都可以让do_sum更快,这将是非常受欢迎的。
为简单起见,我们可以假设lbub不包含nan值,但如果解决方案也能处理它们,那就更好了。
(By顺便说一下,这个代码示例与官方numba文档中说明parallel=True的示例非常相似,但当我尝试在上面的do_sum中添加parallel=True时,我得到了显著的减慢。不知道为什么会这样。)

rdlzhqv9

rdlzhqv91#

我做了很多测试,成功的是并行化

@nb.njit(parallel=True)
def do_sum2(A, lb, ub):
    n = len(A)
    fraction = n // 32
    acc = 0.0
    for j in nb.prange(32):
        acc_local = 0.0
        for i in range(j*fraction, min((j+1)*fraction, n)):
            a = 0.0 if np.isnan(A[i]) else A[i]
            acc += abs(max(min(a, ub[i]), lb[i]))
        acc += acc_local
    return acc

测试代码,如果其他人想尝试优化:

@line_profiler_pycharm.profile
def main():
    N = 10_000_000
    A = np.random.random(N)

    # setting 50% nans
    A[np.random.choice(np.arange(N), N//2, replace=False)] = np.nan
    lb = np.random.random(N)
    ub = lb + np.random.random(N)

    # compiling
    do_sum(A, lb, ub)
    do_sum2(A, lb, ub)

    # test
    result1 = do_sum(A, lb, ub)
    result2 = do_sum2(A, lb, ub)

相关问题