numpy apply_along_axis vectorisation

时间:2017-10-06 04:00:19

标签: python-3.x numpy vectorization

我正在尝试实现一个函数,该函数将每行放在一个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函数?

1 个答案:

答案 0 :(得分:2)

这是一种向量化方法,将zeros设置为NaN,让我们使用np.nanmaxnp.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+

因此,我们看到{{1}}加速。