如何提高 np.random.choice() 循环效率 [英] How to improve np.random.choice() looping efficiency

问题描述

我正在尝试将 np.random.choice 应用于具有不同权重的大数组，想知道有什么方法可以避免循环并提高性能?在这里 len(weights) 可能是数百万.

I am trying to apply np.random.choice to a big array with different weights, and wondering any way could avoid looping and improve the performance? Over here len(weights) could be millions.

weights = [[0.1, 0.5, 0.4],
           [0.2, 0.4, 0.4],
           ...
           [0.3, 0.3, 0.4]]

choice = [1, 2, 3]
ret = np.zeros((len(weights), 20))
for i in range(len(weights)):
    ret[i] = np.random.choice(choice, 20, p=weights[i])

推荐答案

这是我在 随机矩阵的所有行的快速随机加权选择:

def vectorized_choice(p, n, items=None):
    s = p.cumsum(axis=1)
    r = np.random.rand(p.shape[0], n, 1)
    q = np.expand_dims(s, 1) >= r
    k = q.argmax(axis=-1)
    if items is not None:
        k = np.asarray(items)[k]
    return k

p 应该是一个二维数组，其行是概率向量.n 是从每行定义的分布中抽取的样本数.如果items 为None，则样本为range(0, p.shape[1]) 内的整数.如果 items 不是 None，则预期是一个长度为 p.shape[1] 的序列.

p is expected to be a two-dimensional array whose rows are probability vectors. n is the number of samples to draw from the distribution defined by each row. If items is None, the samples are integers in range(0, p.shape[1]). If items is not None, it is expected to be a sequence with length p.shape[1].

示例:

In [258]: p = np.array([[0.1, 0.5, 0.4], [0.75, 0, 0.25], [0, 0, 1], [1/3, 1/3, 1/3]])                                                   

In [259]: p                                                                                                                              
Out[259]: 
array([[0.1       , 0.5       , 0.4       ],
       [0.75      , 0.        , 0.25      ],
       [0.        , 0.        , 1.        ],
       [0.33333333, 0.33333333, 0.33333333]])

In [260]: vectorized_choice(p, 20)                                                                                                       
Out[260]: 
array([[1, 1, 2, 1, 1, 2, 2, 2, 1, 2, 1, 1, 1, 2, 2, 0, 1, 2, 2, 2],
       [0, 2, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 0],
       [2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2],
       [1, 0, 2, 2, 0, 1, 2, 1, 0, 0, 0, 0, 2, 2, 0, 0, 2, 1, 1, 2]])

In [261]: vectorized_choice(p, 20, items=[1, 2, 3])                                                                                      
Out[261]: 
array([[2, 1, 2, 2, 2, 3, 2, 2, 2, 2, 3, 3, 2, 2, 3, 3, 2, 3, 2, 2],
       [1, 1, 1, 1, 1, 1, 1, 1, 3, 3, 3, 1, 1, 3, 3, 1, 3, 1, 1, 1],
       [3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3],
       [3, 3, 3, 1, 3, 2, 1, 2, 3, 1, 2, 2, 3, 2, 1, 2, 1, 2, 2, 2]])

p 形状为 (1000000, 3) 的时间:

In [317]: p = np.random.rand(1000000, 3)

In [318]: p /= p.sum(axis=1, keepdims=True)

In [319]: %timeit vectorized_choice(p, 20, items=np.arange(1, p.shape[1]+1))
1.89 s ± 28.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

以下是 Divakar 的功能时间:

Here's the timing for Divakar's function:

In [320]: %timeit random_choice_prob_vectorized(p, 20, choice=np.arange(1, p.shape[1]+1))
7.33 s ± 43.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

如果你增加p中的列数，差异会不那么明显，如果你让列数足够大，Divakar的函数会更快.例如

The difference will be less pronounced if you increase the number of columns in p, and if you make the number of columns big enough, Divakar's function will be faster. E.g.

In [321]: p = np.random.rand(1000, 120)

In [322]: p /= p.sum(axis=1, keepdims=True)

In [323]: %timeit vectorized_choice(p, 20, items=np.arange(1, p.shape[1]+1))
6.41 ms ± 20.5 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [324]: %timeit random_choice_prob_vectorized(p, 20, choice=np.arange(1, p.shape[1]+1))
6.29 ms ± 342 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

这篇关于如何提高 np.random.choice() 循环效率的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！