用 numba 找到两个向量之间的最小差异答案

【问题标题】：Find minimum difference between two vectors with numba用 numba 找到两个向量之间的最小差异
【发布时间】：2021-06-27 16:27:32
【问题描述】：

我尝试优化搜索两个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)

【问题讨论】：

我可以通过将min_val_idx_a, min_val_idx_b = 0, 0 替换为一个包含两个零值的numpy 数组来实现并行工作。然后我更新了数组中的值。但是，尝试多次运行 timeit 会产生不一致的结果......所以我认为这不是完整的答案。
@jakub 谢谢，对我来说结果是一致的
@jakub 抱歉，有一些错误：0.0 9999 9999，2.2752791872804323e-08 1581 5090，0.0 9999 9999，0.0 9999 9999
numba parallel documentation 有有用的注释。它指出“用户需要确保循环没有交叉迭代依赖项，除了支持的减少。”由于当前的实现具有交叉迭代依赖性，因此可能可以根据那些支持的缩减来重新构建原始问题。
我认为，问题在于您可能会在每次迭代时覆盖索引。这不是减少方法，例如+=。对于执行+= 的循环，numba 可以并行计算多个总和，最后将所有这些总和求和。在当前问题中，索引被覆盖。你看得到差别吗？并行线程（不确定它们是否是线程）如何知道如何减少它们计算的值。似乎是一种竞争条件。

标签： python numpy parallel-processing jit numba

【解决方案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

【讨论】：

我认为您只能像 OP 那样比较矩阵的上半部分。这可能更快，结果应该是相同的，尽管它使实现有点复杂。在实践中，这意味着 b 可以替换为 b[i:] 并且索引会增加以考虑到这种转变。
当您将向量与自身进行比较时，确实如此 [A,B,C]->[AB,AC,BC]。但是在这种情况下，两个向量是不同的，因此有必要遍历所有可能的对 [A,B,C]x[D,E,F]->[AD,AE,AF,BD,BE,BF,CD, CE,CF]

【解决方案2】：

为了完成@aerobiomat 的精彩回答，Numba 目前不支持 显式/推断 最小/最大并行减少（截至日期写作——对于 Numba 0.52）。引用documentation：

二元运算符：+ - * / /? % | >> ^ << & ** // [...]
用户需要确保循环没有交叉迭代依赖，除了支持的缩减。 [...]
如果一个变量被 二元函数/运算符 使用它在循环体中的先前值更新，则自动推断减少。 [...]
编译器可能检测不到这种情况，然后会发生竞态条件。

因此，您的代码格式错误并出现竞争条件。

或者，您可以生成所有值的矩阵并使用隐式并行np.argmin 和parallel=True，但@aerobiomat 提供的解决方案应该更快（因为它不会生成大的临时矩阵）。

【讨论】：