python-3.x - numpy apply_along_axis 矢量化

Question

我正在尝试实现一个函数，该函数获取 numpy 二维数组中的每一行并返回特定计算的标量结果。我当前的代码如下所示:

img = np.array([
    [0,  5,  70, 0,  0,  0 ],
    [10, 50, 4,  4,  2,  0 ],
    [50, 10, 1,  42, 40, 1 ], 
    [10, 0,  0,  6,  85, 64],
    [0,  0,  0,  1,  2,  90]]
)

def get_y(stride):
    stride_vals = stride[stride > 0]
    pix_thresh = stride_vals.max() - 1.5*stride_vals.std()
    return np.argwhere(stride>pix_thresh).mean()

np.apply_along_axis(get_y, 0, img)
>> array([ 2. ,  1. ,  0. ,  2. ,  2.5,  3.5])

它按预期工作，但是，性能并不好，因为在真实数据集中，每帧大约有 2000 行和 20-50 列，每秒出现 60 次。

有没有办法加快这个过程，或许不使用 np.apply_along_axis 函数？

Answer 1

这是一种将 zeros 设置为 NaN 的矢量化方法，让我们使用 np.nanmax 和 np.nanstd 来计算那些 max 和 std 值，避免使用 zeros，就像这样 -

imgn = np.where(img==0, np.nan, img)
mx = np.nanmax(imgn,0) # np.max(img,0) if all are positive numbers
st = np.nanstd(imgn,0)
mask = img > mx - 1.5*st
out = np.arange(mask.shape[0]).dot(mask)/mask.sum(0)

运行时测试-

In [94]: img = np.random.randint(-100,100,(2000,50))

In [95]: %timeit np.apply_along_axis(get_y, 0, img)
100 loops, best of 3: 4.36 ms per loop

In [96]: %%timeit
    ...: imgn = np.where(img==0, np.nan, img)
    ...: mx = np.nanmax(imgn,0)
    ...: st = np.nanstd(imgn,0)
    ...: mask = img > mx - 1.5*st
    ...: out = np.arange(mask.shape[0]).dot(mask)/mask.sum(0)
1000 loops, best of 3: 1.33 ms per loop

因此，我们看到了 3x+ 加速。