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

Question

我试图在一个非常大的二维 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 有序数组没有帮助。