如何在 Rust 中生成伯努利随机整数数组？

Question

我正在尝试创建一个随机 2D 数组，其中每个条目的概率为 0 p，否则为 1。在 python 中，numpy 使其变得简单：

rand_arr = (np.random.rand(nrows, ncols) < p).astype(np.dtype('uint8'))

在 Rust 中，我写了这样的东西

use rand::Rng;
use rand::distributions::Bernoulli;
use ndarray::{Array2};

// Create a random number generator
let mut rng = rand::thread_rng();

// Create a Bernoulli distribution with the given probability
let bernoulli = Bernoulli::new(p).unwrap();

// Generate a 2D array of random boolean values using the Bernoulli distribution
let rand_arr: Array2<bool> = Array2::from_shape_fn((nrows, ncols), |_| {
    rng.sample(bernoulli)
});

// convert bools to 0 and 1
let rand_arr = rand_arr.mapv(|b| if b { 1 } else { 0 });

这非常慢，实际上比 python 版本还慢！这告诉我我做错了什么。我怎样才能编写一个更快的实现？

Answer 1

在开始之前，请允许我将您的代码变成可测试的代码。

use ndarray::Array2;
use rand::distributions::Bernoulli;
use rand::Rng;
use std::hint::black_box;
use std::time::{Duration, Instant};

fn main() {
    let nrows = 200;
    let ncols = 400;
    let p = 0.2;
    let bernoulli = Bernoulli::new(p).unwrap();

    let (rand_arr, t) = original(bernoulli, nrows, ncols);

    println!("{t:?}");
    black_box(rand_arr);
}

fn original(bernoulli: Bernoulli, nrows: usize, ncols: usize) -> (Array2<u8>, Duration) {
    let mut rng = rand::thread_rng();

    let t = Instant::now();
    let rand_arr: Array2<bool> = Array2::from_shape_fn((nrows, ncols), |_| rng.sample(bernoulli));
    let rand_arr = rand_arr.mapv(|b| if b { 1 } else { 0 });

    (rand_arr, t.elapsed())
}

当在我的计算机上以发布模式运行时，此报告大约为 302 微秒。这将是基线。我们现在可以插入不同的函数来对original它们进行计时。

请注意，如果您不指定整数类型，Rust 将使用i32. 我选择u8匹配Python版本。

首先，我将清理您的代码而不改变含义。

fn original_clean(bernoulli: Bernoulli, nrows: usize, ncols: usize) -> (Array2<u8>, Duration) {
    let rng = rand::thread_rng();
    let mut iter = rng.sample_iter(bernoulli);

    let t = Instant::now();
    let rand_arr =
        Array2::from_shape_simple_fn((nrows, ncols), || if iter.next().unwrap() { 1 } else { 0 });

    (rand_arr, t.elapsed())
}

这比多次sample_iter调用效率更高。由于您不使用索引，因此sample它也会使用。from_shape_simple_fn但最重要的是，它就地进行从布尔值到整数的映射。

这一过程需要 237 微秒，这一时间相当长，但了解数组创建和随机生成之间的时间分配方式会很有用。

fn not_rand(_bernoulli: Bernoulli, nrows: usize, ncols: usize) -> (Array2<u8>, Duration) {
    let t = Instant::now();
    let rand_arr = Array2::from_shape_simple_fn((nrows, ncols), || 1);

    (rand_arr, t.elapsed())
}

这会生成一个全为 1 的数组。需要 24 微秒。所以90%的时间都花在了随机生成上。我去了解一下RNG NumPy用的是什么。

default_rng目前用作PCG64默认值BitGenerator

让我们看看rand 使用什么。

目前使用的算法是12轮的ChaCha分组密码。

我不太熟悉不同 RNG 的性能，但 ChaCha 是一个加密安全的 RNG，所以它可能会更慢。rust 中有一个 PCG-64 的箱子：rand_pcg。我用它做了一个函数来插入计时器。

fn pcg64(bernoulli: Bernoulli, nrows: usize, ncols: usize) -> (Array2<u8>, Duration) {
    use rand::SeedableRng;
    let rng = rand_pcg::Pcg64::from_entropy();
    let mut iter = rng.sample_iter(bernoulli);

    let t = Instant::now();
    let rand_arr =
        Array2::from_shape_simple_fn((nrows, ncols), || if iter.next().unwrap() { 1 } else { 0 });

    (rand_arr, t.elapsed())
}

运行时间为 132 微秒！如果您的用例允许，您甚至可以使用更便宜的 RNG（Rust 和 Python）。

如果这仍然比 NumPy 慢，我就不会那么惊讶。NumPy 的性能密集型代码是用 C 编写的，因此它不会受到 Python 通常缓慢的影响。然而，如果 NumPy 外部的逻辑花费了大量时间，Rust 版本应该击败 Python。

这是操场上的所有内容，除了 PCG，因为板条箱不可用。

我继续测试 Python 只是为了检查一下。它的计时并不完全相同，但它应该可以让您了解 Rust 版本的表现。

编辑：我没有注意到它random.rand使用旧的 MT19937 RNG，所以我已经更新它以通过default_rng. MT19937 大约慢 60%。我也尝试过Rust的mt19937，但速度相当慢。

import numpy as np
import time

nrows = 200
ncols = 400
p = 0.2

a = time.perf_counter()
rand_arr = (np.random.default_rng().random((nrows, ncols), dtype='float32') < p).astype(np.dtype('uint8'))
b = time.perf_counter()

print((b - a) * 1000 * 1000, "microseconds")

这需要 370 微秒。