python - 用numba找到两个向量之间的最小差异

Question

我尝试使用 numpy 优化搜索两个 numba 向量之间的最小值。在我使用 prange 和 parallel=True 选项之前，有加速并且结果是正确的。我知道问题在于在并行执行期间共享变量 min_val 、 tmp 、 min_val_idx_a 、 min_val_idx_b (可能使用并行线程)。有没有办法克服这个问题并在 numba 选项中使用 parallel=True ？ (这使我的简单代码快了 300 倍)

import numpy as np
from numba import jit, int32, void, double, prange

@jit(void(double[:], double[:], int32), nopython=True, parallel=True)
def lowest_value_numba(a, b, n): 
    
    # initialization
    min_val_idx_a, min_val_idx_b = 0, 0
    min_val = tmp = np.abs(a[0]-b[0])
    
    for i in prange(n):
#         print(i)
        for j in prange(i, n):
            tmp = np.abs(a[i]-b[j])
            
            if(tmp < min_val):
                min_val = tmp            
                min_val_idx_a = i
                min_val_idx_b = j
    print(min_val, min_val_idx_a, min_val_idx_b)
                
n = int(1e4)
a = np.random.uniform(low=0.0, high=1.0, size=n)
b = np.random.uniform(low=0.0, high=1.0, size=n)

# setting min value by setting the same valu efor a[n-1] and b[n-1] 
a[n-1], b[n-1] = 1, 1

%timeit -n 1 -r 1 lowest_value_numba(a, b, n)

输出不正确(应该是 0.0 9999 9999 ):

0.23648058275546968 0 0
223 µs ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)

但是对于使用 parallel=False 输出的编译是正确的(最后一个值彼此最接近):

0.0 9999 9999
65 ms ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)

Answer 1

如果按行并行化，则可以避免有关交叉迭代依赖项的评论问题。例如:

from numba import jit, int32, double, prange
from numba.types import Tuple

@jit(Tuple([double, int32, int32])(double[:], double[:], int32),
     nopython=True, parallel=True)
def lowest_value_numba(a, b, n):

    min_dif = np.empty_like(a)
    min_j = np.empty((n,), dtype=int32)

    for i in prange(n):
        diff = np.abs(a[i] - b)
        min_j[i] = j = np.argmin(diff)
        min_dif[i] = diff[j]

    i = np.argmin(min_dif)
    j = min_j[i]
    min_val = min_dif[i]

    return min_val, i, j

结果与您的实现(使用 parallel=False 和 for j in prange(n) 测试)和蛮力 Numpy 方法一致:

def lowest_value_numpy(a, b):
    diff = np.abs(np.atleast_2d(a).T - np.atleast_2d(b))
    indices = np.unravel_index(diff.argmin(), diff.shape)
    return diff[indices], *indices