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

2ul0zpep 于 2023-03-02 发布在其他

关注(0)|答案(1)|浏览(114)

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

numpy

来源：https://stackoverflow.com/questions/75585818/numba-is-it-possible-to-do-absolute-sum-without-nan-faster

1条答案

按热度按时间

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)

赞(0）回复(0）举报 2023-03-02

我来回答

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

1条答案

相关问题

热门标签

最新问答