生成包含随机布尔值的大型 numpy 数组的内存有效方法

Question

我需要创建一个包含随机布尔值的大型 numpy 数组，而不会点击交换。

我的笔记本电脑有 8 GB 的 RAM。创建一个(1200, 2e6) 数组需要不到 2 秒的时间并使用 2.29 GB 的 RAM：

>>> dd = np.ones((1200, int(2e6)), dtype=bool)
>>> dd.nbytes/1024./1024
2288.818359375

>>> dd.shape
(1200, 2000000)

对于相对较小的(1200, 400e3)，np.random.randint 仍然相当快，大约需要 5 秒来生成一个 458 MB 的数组：

db = np.array(np.random.randint(2, size=(int(400e3), 1200)), dtype=bool)
print db.nbytes/1024./1024., 'Mb'

但是，如果我将数组的大小加倍到 (1200, 800e3)，我会点击交换，创建 db 需要大约 2.7 分钟；(

cmd = """
import numpy as np
db = np.array(np.random.randint(2, size=(int(800e3), 1200)), dtype=bool)
print db.nbytes/1024./1024., 'Mb'"""

print timeit.Timer(cmd).timeit(1)

使用random.getrandbits 需要更长的时间（~8 分钟），并且还使用交换：

from random import getrandbits
db = np.array([not getrandbits(1) for x in xrange(int(1200*800e3))], dtype=bool)

将np.random.randint 用于(1200, 2e6) 只会给出MemoryError。

有没有更有效的方法来创建(1200, 2e6) 随机布尔数组？

Answer 1

使用np.random.randint 的一个问题是它生成64 位整数，而numpy 的np.bool dtype 仅使用8 位来表示每个布尔值。因此，您分配的中间数组比需要的大 8 倍。

避免中间 64 位 dtypes 的解决方法是使用 np.random.bytes 生成随机字节字符串，可以使用 np.fromstring 将其转换为 8 位整数数组。然后可以将这些整数转换为布尔值，例如通过测试它们是否小于 255 * p，其中 p 是每个元素为True 的期望概率：

import numpy as np

def random_bool(shape, p=0.5):
    n = np.prod(shape)
    x = np.fromstring(np.random.bytes(n), np.uint8, n)
    return (x < 255 * p).reshape(shape)

基准测试：

In [1]: shape = 1200, int(2E6)

In [2]: %timeit random_bool(shape)
1 loops, best of 3: 12.7 s per loop

一个重要的警告是，概率将向下舍入到最接近的 1/256 倍数（对于 1/256 的精确倍数，例如 p=1/2，这不会影响准确性）。

---

更新：

一种更快的方法是利用这样一个事实，即您只需要在输出数组中为每个 0 或 1 生成一个随机位。因此，您可以创建一个 8 位整数的随机数组，其大小为最终输出的 1/8，然后使用 np.unpackbits 将其转换为 np.bool：

def fast_random_bool(shape):
    n = np.prod(shape)
    nb = -(-n // 8)     # ceiling division
    b = np.fromstring(np.random.bytes(nb), np.uint8, nb)
    return np.unpackbits(b)[:n].reshape(shape).view(np.bool)

例如：

In [3]: %timeit fast_random_bool(shape)
1 loops, best of 3: 5.54 s per loop