为什么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))

- Blink

---

这并不是一个公平的比较。每次，numpy都必须对p列表进行累加求和，并将其放入新的向量中，然后迭代它。你实际上是通过知道只有三个变量，并且第一个和第三个变量的和为0.5来进行预处理。除此之外，正如下面所指出的那样，numpy被优化为向量化操作，而不是执行数千次单个简单操作的工具。 - David Robinson

此外，请使用 timeit 而不是单独使用 time。 - Marcin

请参考 non-repetitive-random-number-in-numpy 获取在 Numpy 1.8.1 中选择 10000^2 中的 40000 个不重复随机数的方法。 - denis

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

- user2357112

---

你所比较的是同样的东西吗？第一个计算了10^3 * 10^4 = 10^7个随机数，但第二个计算了10^2 * 10^3 * 10^4 = 10^9个随机数，对吧？ - DSM

我可能忘记在8年半前回复，或者回复在某个时候被删除了，但是DSM指出的问题已经得到解决（在DSM发布后的3分钟内）。 - user2357112

Answer 2

这个使用累积分数的解决方案快了25倍：

def choice(options,probs):
    x = np.random.rand()
    cum = 0
    for i,p in enumerate(probs):
        cum += p
        if x < cum:
            break
    return options[i]


options = ['a','b','c','d']
probs = [0.2,0.6,0.15,0.05]
runs = 100000


now = time.time()
temp = []
for i in range(runs):
    op = choice(options,probs)
    temp.append(op)
temp = Counter(temp)
for op,x in temp.items():
    print(op,x/runs)
print(time.time()-now)

print("")
now = time.time()
temp = []
for i in range(runs):
    op = np.random.choice(options,p = probs)
    temp.append(op)
temp = Counter(temp)
for op,x in temp.items():
    print(op,x/runs)
print(time.time()-now)

运行它，我得到：

b 0.59891
a 0.20121
c 0.15007
d 0.04981
0.16232800483703613

b 0.5996
a 0.20138
c 0.14856
d 0.05046
3.8451428413391113

- Miguel

---

小重构： def fast_choice(options, probs): x = random.random()#np.random.rand() cum = 0 for i, p in enumerate(probs): cum += p if x < cum: return options[i] return options[-1] - Demetry Pascal

你的列表不是一个numpy数组。这会给numpy带来很大的开销。 - Fırat Kıyak

Answer 3

我花了很长时间才发现，由于使用 np.random.choice 进行随机键抽样，我的数据生成器非常慢。

如果不需要非均匀分布，则可以使用下面这个可行的解决方案。

替换

def get_random_key(a_huge_key_list):
    return np.random.choice(a_huge_key_list)

使用

def get_random_key(a_huge_key_list):
    L = len(a_huge_key_list)
    i = np.random.randint(0, L)
    return a_huge_key_list[i]

这将使速度提高60倍。

- pitfall

---

嗯，有趣。你知道为什么这会对速度产生如此大的影响吗？从源代码来看，你扩展的版本看起来与NumPy已经在做的事情并没有太大的区别。 - Mark Dickinson

经过一些测试，我怀疑你看到的时间差异的主要原因是你使用了一个实际的Python list作为输入，而np.random.choice隐式地将其转换为NumPy数组，而你的第二个版本的get_random_key避免了这种转换。对于大型列表，列表到数组的转换成为瓶颈。当我在一个包含10 ** 6个元素的1d NumPy _array_上测试这两个变量时，我的计时结果更接近：第二个版本大约快50％。 - Mark Dickinson

Answer 4

对于那些偶然遇到这个问题并且不是抽取1000个样本10000次而是抽取1个样本10000次的人，自Python 3.6以来存在一种更快的替代方法。函数random.choices相比numpy.random.choice快大约20倍。

timeit("random.choices([-1, 0, 1], k=1, weights=[.25, .5, .25])", 
setup="import random", number=10000)
# >>> 0.018841781999981322

vs

timeit("numpy.random.choice([-1, 0, 1], size=1, p=[.25, .5, .25])", 
setup="import numpy", number=10000)
# >>> 0.40612822200000664

- Dennis Proksch

---

哇，这差别太大了。谢谢！ - Mdev

Answer 5

我怀疑np.random.choice的普适性会使其减慢速度，特别是对于小样本而言更为明显。

对于if版本的粗略向量化如下：

def foo(n):
    x = np.random.rand(n)
    var = np.zeros(n)
    var[x<.25] = -1
    var[x>.75] = 1
    return var

在使用 ipython 时，我遇到了以下问题：

timeit np.random.choice([-1,0,1],size=1000,p=[.25,.5,.25])
1000 loops, best of 3: 293 us per loop

timeit foo(1000)
10000 loops, best of 3: 83.4 us per loop

timeit np.random.choice([-1,0,1],size=100000,p=[.25,.5,.25])
100 loops, best of 3: 11 ms per loop

timeit foo(100000)
100 loops, best of 3: 8.12 ms per loop

所以对于大小为1000的情况，choice会慢3-4倍，但对于更大的向量，差异开始消失。

- hpaulj

Answer 6

其他答案至少涉及以下一项：

1- 使用Python列表作为numpy.random.choice的输入并创建开销。

2- 使用先前的知识，即len（array）将为3。

3- 分布是均匀的。

对于任意长度的列表，最快的算法之一是在每个步骤中将列表分成两部分。例如，以下代码将适用于通用情况。

def my_random_function(collection, p):
    miles = []
    current = 0
    for prob in p:
        miles.append(current)
        current += prob
    if not math.isclose(current,1):
        raise ValueError()
    x = random.random()
    _all = list(zip(collection,miles))
    while(len(_all)!= 1):
        if _all[len(_all)//2][1] < x:
            _all = _all[len(_all)//2:]
        else:
            _all = _all[0: len(_all)//2]
    return _all[0][0]

为了比较差异，我准备了两种情况：

small_list = list(range(3))
small_array = np.arange(3)
#create a random probability list
small_p = [random.random() for i in range(3)]
small_p = [prob/sum(small_p) for prob in small_p]
small_p_np = np.array(small_p)

large_list = list(range(10000))
large_array = np.arange(10000)
#create a random probability list
large_p = [random.random() for i in range(10000)]
large_p = [prob/sum(large_p) for prob in large_p]
large_p_np = np.array(large_p)

结果如下：

%timeit np.random.choice(small_array, p= small_p_np)

68.1 µs ± 196 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%timeit my_random_function(small_list, small_p)

5.13 µs ± 26.9 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

%timeit np.random.choice(large_array, p= large_p_np)

279 µs ± 1.15 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit my_random_function(large_list, large_p)

3.26 ms ± 5.82 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

根据结果显示，对于小的集合，numpy.random.choice的执行时间超过x10，但是当元素数量增加时，它很快成为更好的选择。似乎该函数有很大的开销，在代码性能关键部分的小列表中最好避免使用。

- Fırat Kıyak