NumPy apply_along_axis 向量化

Question

我正在尝试实现一个函数，该函数接受numpy 2d数组中的每一行并返回某些计算的标量结果。我的当前代码如下:

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])

功能正常，但是性能并不好，因为在真实数据集中，每个框架大约有2k行和20-50列，每秒出现60次。

是否有一种方法可以加快处理速度，比如不使用np.apply_along_axis函数？

- ymoiseev

---

我预计 apply_along_axis 比简单的行迭代要慢。 - hpaulj

@hpaulj我测试了一下，发现按行迭代的速度确实稍慢。有什么特殊的原因吗？如果沿轴应用通常比行迭代更慢，那么拥有这样的功能有什么意义呢？ - ymoiseev

方便。这是Python代码，您可以阅读它。它不会编译循环或函数的评估。它只是将循环泛化，使得在多个维度（即4d数组的另外3个维度）上表达评估更容易。它使用np.ndindex生成循环索引。 - hpaulj

Answer 1

这里有一种向量化的方法，将 zeros 设置为 NaN，然后我们可以使用 np.nanmax 和 np.nanstd 来计算那些避免 zeros 的 max 和 std 值，代码如下 -

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+的加速。

- Divakar

---

谢谢，这个很好用。对于更大的图像，它可以提供100倍以上的加速！ - ymoiseev