np.nan 和 np.inf 的 Numba 性能问题答案

【问题标题】：Numba performance issue with np.nan and np.infnp.nan 和 np.inf 的 Numba 性能问题
【发布时间】：2022-01-17 09:52:56
【问题描述】：

我正在使用numba 来加速我的代码。我注意到在函数内部使用np.inf 而不是np.nan 时，性能会有很大差异。下面我附上了三个示例函数进行说明。

function1 不会被 numba 加速。
function2 和 function3 都由 numba 加速，但一个使用 np.nan 而另一个使用 np.inf。

在我的机器上，三个函数的平均运行时间分别为0.032284s、0.041548s和0.019712s。使用np.nan 似乎比np.inf 慢得多。为什么性能差异很大？提前致谢。

编辑：我正在使用Python 3.7.11 和Numba 0.55.Orc1。

import numpy as np
import numba as nb

def function1(array1, array2):
    nr, nc = array1.shape
    output1 = np.empty((nr, nc), dtype='float')
    output2 = np.empty((nr, nc), dtype='float')
    output1[:] = np.nan
    output2[:] = np.nan

    for r in range(nr):
        row1 = array1[r]
        row2 = array2[r]
        diff = row1 - row2
        id_threshold =np.nonzero( (row1 - row2) > 8 )
        output1[r][id_threshold] = 1
        output2[r][id_threshold] = 0

    output1 = output1.flatten()
    output2 = output2.flatten()
    id_keep = np.nonzero(output1 != np.nan)
    output1 = output1[id_keep]
    output2 = output2[id_keep]
    output = np.vstack((output1, output2))
    return output

@nb.njit('float64[:,::1](float64[:,::1], float64[:,::1])', parallel=True)
def function2(array1, array2):
    nr, nc = array1.shape
    output1 = np.empty((nr,nc), dtype='float')
    output2 = np.empty((nr, nc), dtype='float')
    output1[:] = np.nan
    output2[:] = np.nan

    for r in nb.prange(nr):
        row1 = array1[r]
        row2 = array2[r]
        diff = row1 - row2
        id_threshold =np.nonzero( (row1 - row2) > 8 )
        output1[r][id_threshold] = 1
        output2[r][id_threshold] = 0

    output1 = output1.flatten()
    output2 = output2.flatten()
    id_keep = np.nonzero(output1 != np.nan)
    output1 = output1[id_keep]
    output2 = output2[id_keep]
    output = np.vstack((output1, output2))
    return output

@nb.njit('float64[:,::1](float64[:,::1], float64[:,::1])', parallel=True)
def function3(array1, array2):
    nr, nc = array1.shape
    output1 = np.empty((nr,nc), dtype='float')
    output2 = np.empty((nr, nc), dtype='float')
    output1[:] = np.inf
    output2[:] = np.inf

    for r in nb.prange(nr):
        row1 = array1[r]
        row2 = array2[r]
        diff = row1 - row2
        id_threshold =np.nonzero( (row1 - row2) > 8 )
        output1[r][id_threshold] = 1
        output2[r][id_threshold] = 0
    output1 = output1.flatten()
    output2 = output2.flatten()
    id_keep = np.nonzero(output1 != np.inf)
    output1 = output1[id_keep]
    output2 = output2[id_keep]
    output = np.vstack((output1, output2))
    return output


array1 = 10*np.random.random((1000,1000))
array2 = 10*np.random.random((1000,1000))

output1 = function1(array1, array2)
output2 = function2(array1, array2)
output3 = function3(array1, array2)

【问题讨论】：

您使用的是哪个版本的 numba？我在0.51.2 上收到错误消息。
@hilberts_drinking_problem 请查看编辑后的帖子。

标签： python numpy performance numba

【解决方案1】：

第二个要慢得多，因为output1 != np.nan 返回一个副本output1，因为np.nan != np.nan 是True（与任何其他值一样——v != np.nan 始终为真）。因此，要计算的结果数组要大得多，从而导致执行速度变慢。

关键是您必须切勿使用比较运算符将值与np.nan 进行比较：改用np.isnan(value)。在您的情况下，您应该使用np.logical_not(np.isnan(output1))。

由于np.logical_not 创建的临时数组，第二个实现可能会稍微慢一些（我没有看到在我的机器上使用 NaN 或 Inf 在代码被更正后有任何统计学上的显着差异）。

【讨论】：