Python Numba 非确定性结果

Question

我正在尝试使用 numba.njit 来加速一些数学函数。

该函数接受一个数字数组作为输入，并且对于每个元素 i，它检查有多少邻居（在某个定义的邻近度 n 中）小于 i。这被称为 局部等级变换 并且此函数是 100% 确定性的。

所以我写了这个并且它的行为符合预期：

import numba as nb
import numpy as np


@nb.njit(["float64[:](float64[:], int32)",
          "float32[:](float32[:], int32)",
          "uint8[:](uint8[:], int32)"], parallel=True, nogil=True)
def local_rank_transform_jit(sig, n):
    pad = n//2
    result = np.zeros_like(sig)
    padded_sig = np.zeros(len(sig) + n - 1)
    padded_sig[pad:-pad] = sig
    
    for idx in nb.prange(pad, len(result) + pad):
        neighb = padded_sig[idx - pad:idx + pad]
        result[idx - pad] = np.sum(neighb < padded_sig[idx])
    
    return result

然后我意识到我将运行对许多输入进行此操作并且我可以将其并行化，所以我编写了以下函数：


@nb.njit(["float64[:,:](float64[:,:], int32)",
          "float32[:,:](float32[:,:], int32)",
          "uint8[:,:](uint8[:,:], int32)"], parallel=True, nogil=True)
def local_rank_transform_many_jit(signals, n):
    pad = n//2
    result = np.zeros_like(signals[0])
    all_results = np.zeros_like(signals)
    padded_sig = np.zeros(len(result) + n - 1)
    
    for i in nb.prange(len(signals)):
            padded_sig[pad:-pad] = signals[i]

            for j in nb.prange(pad, len(result) + pad):
                neighb = padded_sig[j - pad:j + pad + 1]
                result[j - pad] = np.sum(neighb < padded_sig[j])
            
            all_results[i] = result
            
    return all_results

但是当使用这个时，我得到了奇怪的结果。

首先，结果是不确定的，即我运行函数两次使用相同的输入并得到不同的输出。

然后，我还注意到我得到了 错误的 结果，例如：

> signals
array([[2603., 1352., 2087.,  240.,  979.],
       [1633., 1181., 1328., 1956., 2808.],
       [1628., 2233., 1781.,  448.,  737.],
       [ 574., 2332., 1245.,  246., 2697.],
       [2170.,  312.,  942.,  811., 1497.]])


> local_rank_transform_many_jit(signals, 3)
array([[3., 0., 1., 1., 2.],
       [3., 0., 1., 1., 2.],
       [3., 0., 1., 1., 2.],
       [1., 2., 1., 0., 2.],
       [3., 0., 1., 1., 2.]])

在这种情况下，3 是不可能的输出，因为一个元素在 3 的邻域中不能大于 3 个元素。这意味着它大于自身！

我怀疑并行化造成了某种竞争条件，但我不知道我做错了什么。

Answer 1

padded_sig[pad:-pad] 被不同的值 signals[i] 并行覆盖，因此

的结果

neighb = padded_sig[j - pad:j + pad + 1]

和

result[j - pad] = np.sum(neighb < padded_sig[j])

不确定。

您需要为每个 i 单独的数组 padded_sig，或者确保为每个 i 写入一个独特的 padded_sig 片段。

Answer 2

实施mkrieger1接受的答案后的工作代码：

@nb.njit(["float64[:,:](float64[:,:], int32)",
          "float32[:,:](float32[:,:], int32)",
          "uint8[:,:](uint8[:,:], int32)"], parallel=True, nogil=True)
def local_rank_transform_many_jit(signals, n):
    pad = n//2
    num_sig, sig_len = signals.shape
    all_results = np.zeros_like(signals)
    padded_signals = np.zeros((num_sig, sig_len + n - 1))
    padded_signals[:, pad:-pad] = signals
    
    for i in nb.prange(num_sig):
        for j in nb.prange(pad, sig_len + pad):
            neighb = padded_signals[i, j - pad:j + pad + 1]
            all_results[i, j - pad] = np.sum(neighb < padded_signals[i, j])

    return all_results

Python Numba 非确定性结果

Python Numba non deterministic results

python-3.x

numpy

parallel-processing

numba