python - 生成 numpy 数组或随机分布的 0 和 1 的最快方法

Question

我需要为特定神经网络的 dropout 生成掩码。我正在寻找使用 numpy(仅限 CPU)实现这一目标的最快方法。

我试过:

def gen_mask_1(size, p=0.75):
    return np.random.binomial(1, p, size)


def gen_mask_2(size, p=0.75):
    mask = np.random.rand(size)
    mask[mask>p]=0
    mask[mask!=0]=1
    return mask

其中 p 是 1

的概率

这两种方法的速度相当。

%timeit gen_mask_1(size=2048)
%timeit gen_mask_2(size=2048)

45.9 µs ± 575 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
47.4 µs ± 372 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

有没有更快的方法？

更新

根据目前得到的建议，我测试了一些额外的实现。设置 parallel=True 时，我无法让 @njit 工作(TypingError:在 nopython 模式管道中失败(步骤:转换为 parfors) )，它可以在没有但我认为效率较低的情况下工作。我在此处找到了英特尔 mlk_random 的 python 绑定(bind)(感谢@MatthieuBrucher 的提示!):https://github.com/IntelPython/mkl_random到目前为止，将 mlk_random 与@nxpnsv 的方法结合使用可获得最佳结果。

@njit
def gen_mask_3(size, p=0.75):
    mask = np.random.rand(size)
    mask[mask>p]=0
    mask[mask!=0]=1
    return mask

def gen_mask_4(size, p=0.75):
    return (np.random.rand(size) < p).astype(int)

def gen_mask_5(size):
    return np.random.choice([0, 1, 1, 1], size=size)

def gen_mask_6(size, p=0.75):
    return (mkl_random.rand(size) < p).astype(int)

def gen_mask_7(size):
    return mkl_random.choice([0, 1, 1, 1], size=size)

%timeit gen_mask_4(size=2048)
%timeit gen_mask_5(size=2048)
%timeit gen_mask_6(size=2048)
%timeit gen_mask_7(size=2048)

22.2 µs ± 145 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
25.8 µs ± 336 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
7.64 µs ± 133 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
29.6 µs ± 1.18 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

Answer 1

您可以使用 Numba编译器，并通过在您的函数上应用 njit 装饰器来加快速度。下面是一个非常大的 size

示例

from numba import njit

def gen_mask_1(size, p=0.75):
    return np.random.binomial(1, p, size)

@njit(parallel=True)
def gen_mask_2(size, p=0.75):
    mask = np.random.rand(size)
    mask[mask>p]=0
    mask[mask!=0]=1
    return mask

%timeit gen_mask_1(size=100000)
%timeit gen_mask_2(size=100000)

---

2.33 ms ± 215 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
512 µs ± 25.1 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)