python - 为什么 numpy.random.choice 这么慢？

Question

在编写脚本时，我发现了 numpy.random.choice 函数。我实现它是因为它比等效的 if 语句干净得多。然而，在运行脚本后我意识到它比 if 语句慢明显。

下面是一个MWE。第一种方法需要 0.0 秒，而第二种方法需要 7.2 秒。如果你扩大 i 循环，你会看到 random.choice 变慢的速度有多快。

谁能评论一下为什么 random.choice 这么慢？

import numpy as np
import numpy.random as rand
import time as tm

#-------------------------------------------------------------------------------

tStart = tm.time()
for i in xrange(100):
    for j in xrange(1000):
        tmp = rand.rand()
        if tmp < 0.25:
            var = 1
        elif tmp < 0.5:
            var = -1
print('Time: %.1f s' %(tm.time() - tStart))

#-------------------------------------------------------------------------------

tStart = tm.time()
for i in xrange(100):
    for j in xrange(1000):
        var = rand.choice([-1, 0, 1], p = [0.25, 0.5, 0.25])
print('Time: %.1f s' %(tm.time() - tStart))

Answer 1

你用错了。向量化操作，否则 numpy 将没有任何好处:

var = numpy.random.choice([-1, 0, 1], size=1000, p=[0.25, 0.5, 0.25])

时序数据:

>>> timeit.timeit('''numpy.random.choice([-1, 0, 1],
...                                      size=1000,
...                                      p=[0.25, 0.5, 0.25])''',
...               'import numpy', number=10000)
2.380380242513752

>>> timeit.timeit('''
... var = []
... for i in xrange(1000):
...     tmp = rand.rand()
...     if tmp < 0.25:
...         var.append(1)
...     elif tmp < 0.5:
...         var.append(-1)
...     else:
...         var.append(0)''',
... setup='import numpy.random as rand', number=10000)
5.673041396894519