有没有办法让numpy.argmin()和min()一样快？

Question

我试图找到一个非常大的2D numpy数组的一个维度上的最小数组索引.我发现这很慢(已经尝试过加速瓶颈,这只是一个很小的改进).但是,采用直线最小值似乎要快一个数量级:

import numpy as np
import time

randvals = np.random.rand(3000,160000)
start = time.time()
minval = randvals.min(axis=0)
print "Took {0:.2f} seconds to compute min".format(time.time()-start)
start = time.time()
minindex = np.argmin(randvals,axis=0)
print "Took {0:.2f} seconds to compute argmin".format(time.time()-start)

在我的机器上输出:

Took 0.83 seconds to compute min
Took 9.58 seconds to compute argmin

有没有理由说argmin这么慢？有没有什么方法可以加快到与min相当的速度？

Answer 1

In [1]: import numpy as np

In [2]: a = np.random.rand(3000, 16000)

In [3]: %timeit a.min(axis=0)
1 loops, best of 3: 421 ms per loop

In [4]: %timeit a.argmin(axis=0)
1 loops, best of 3: 1.95 s per loop

In [5]: %timeit a.min(axis=1)
1 loops, best of 3: 302 ms per loop

In [6]: %timeit a.argmin(axis=1)
1 loops, best of 3: 303 ms per loop

In [7]: %timeit a.T.argmin(axis=1)
1 loops, best of 3: 1.78 s per loop

In [8]: %timeit np.asfortranarray(a).argmin(axis=0)
1 loops, best of 3: 1.97 s per loop

In [9]: b = np.asfortranarray(a)

In [10]: %timeit b.argmin(axis=0)
1 loops, best of 3: 329 ms per loop

也许min是足够聪明,可以在数组上顺序完成工作(因此具有缓存局部性),并argmin在数组中跳转(导致大量缓存未命中)？

无论如何,如果你愿意randvals从一开始就保留作为Fortran排序的数组,它会更快,尽管复制到Fortran-ordered也无济于事.

Answer 2

我只是看了一下源代码,虽然我不完全理解为什么事情按照它们的方式完成,但这是发生的事情:

1. np.min基本上是一个电话np.minimum.reduce.

2. np.argmin 首先将您想要操作的轴移动到形状元组的末尾,然后使其成为一个连续的数组,这当然会触发完整数组的副本,除非轴是最后一个开始的轴.

由于正在制作副本,您可以获得创意并尝试实例化更便宜的数组:

a = np.random.rand(1000, 2000)

def fast_argmin_axis_0(a):
    matches = np.nonzero((a == np.min(a, axis=0)).ravel())[0]
    rows, cols = np.unravel_index(matches, a.shape)
    argmin_array = np.empty(a.shape[1], dtype=np.intp)
    argmin_array[cols] = rows
    return argmin_array

In [8]: np.argmin(a, axis=0)
Out[8]: array([230, 532, 815, ..., 670, 702, 989], dtype=int64)

In [9]: fast_argmin_axis_0(a)
Out[9]: array([230, 532, 815, ..., 670, 702, 989], dtype=int64)

In [10]: %timeit np.argmin(a, axis=0)
10 loops, best of 3: 27.3 ms per loop

In [11]: %timeit fast_argmin_axis_0(a)
100 loops, best of 3: 15 ms per loop

我不会把当前的实现称为bug,因为numpy可能有很好的理由做它做的事情,但是这种技巧可以加速应该是高度优化的函数,强烈建议事情可以做得更好.