如何在Python中加速填充numpy数组？

Question

我将尝试使用以下代码填充预分配的字节数组：

```python

我正在尝试使用以下代码填充预分配的字节数组：

```

# preallocate a block array
dt = numpy.dtype('u8')
in_memory_blocks = numpy.zeros(_AVAIL_IN_MEMORY_BLOCKS, dt)

...

# write all the blocks out, flushing only as desired
blocks_per_flush_xrange = xrange(0, blocks_per_flush)
for _ in xrange(0, num_flushes):
    for block_index in blocks_per_flush_xrange:
        in_memory_blocks[block_index] = random.randint(0, _BLOCK_MAX)

    print('flushing bytes stored in memory...')

    # commented out for SO; exists in actual code
    # removing this doesn't make an order-of-magnitude difference in time
    # m.update(in_memory_blocks[:blocks_per_flush])

    in_memory_blocks[:blocks_per_flush].tofile(f)

一些要点:

• num_flushes 很低，大约在 4-10 之间
• blocks_per_flush 值很大，大约是百万级别的
• in_memory_blocks 可以是相当大的缓冲区（我将其设置为1MB到100MB），但时间非常稳定...
• _BLOCK_MAX 是一个8字节无符号整数的最大值
• m 是一个 hashilib.md5()

使用上述代码生成1MB需要约1秒钟; 生成500MB需要约376秒。 相比之下，我的简单 C 程序使用 rand() 可以在 8 秒钟内创建一个 500MB 文件。

如何提高上面 loop 的性能？ 我很确定我忽略了一些明显的东西，导致运行时巨大的差异。

- Allen George

---

使用迭代器遍历任意类型（Python）的效率与编译后的特定类型迭代器（Numpy内部迭代器）相比极其缓慢。如果您想要一个随机整数数组，就像您在示例中所做的那样，请使用numpy.random.randint函数。但我怀疑这只是为了说明目的。您需要发布您实际填充数组的内容以获得更具体的帮助。 - Paul

不，我确实正在创建随机数据并计算其校验和。我认为内存分配是瓶颈，而不是迭代。我可以想象内部的numpy迭代器会更快（因为底层的C实现将使用指针算术）。 - Allen George

请澄清：_BLOCK_MAX 是否等于 255 或 9223372036854775807，以及 dt 是否等于 numpy.uint8 或 numpy.uint64？ - jfs

澄清一下：我的代码中，dt 是 numpy.uint64 类型，而 _BLOCK_MAX 是 uint64 的上限。 - Allen George

Answer 1

由于0.._BLOCK_MAX覆盖了numpy.uint8的所有可能值（我假设numpy.dtype('u8')（即numpy.uint64是一个笔误），你可以使用：

import numpy as np

for _ in xrange(0, num_flushes):
    in_memory_blocks = np.frombuffer(np.random.bytes(blocks_per_flush),
                                     dtype=np.uint8)

    print('flushing bytes stored in memory...')
    # ...

这个版本比@hgomersall的版本快了约8倍：

$ python -mtimeit -s'import numpy as np' '
>     np.uint8(np.random.randint(0,256,20000000))'
10 loops, best of 3: 316 msec per loop

$ python -mtimeit -s'import numpy as np' '
>     np.frombuffer(np.random.bytes(20000000), dtype=np.uint8)'
10 loops, best of 3: 38.6 msec per loop

如果 numpy.dtype('u8') 不是打错了，而你确实需要 numpy.uint64，那么：

a = np.int64(np.random.random_integers(0, _BLOCK_MAX, blocks_per_flush))
in_memory_blocks = a.view(np.uint64) # unsigned

注意：如果数组的数据类型已经是np.int64，则np.int64（）不会进行复制。 .view(numpy.uint64)将其解释为无符号数（也不执行复制）。

- jfs

---

我有一个要求，需要至少执行 n 次磁盘写入操作。 - Allen George

使用 np.random.randint() 的 #1，实际上会返回比 #2 多4到8倍的字节数（取决于您机器上默认的 int 类型）。例如，我可以天真地想象通过生成 int(num_bytes_requested/unsigned_int_size) 个随机整数并填充剩余的字节来实现 np.random.bytes。 这可能可以解释性能差异。 - Allen George

@Allen George：我只在dtype=np.uint8的情况下使用np.random.bytes()，是的，在这种情况下，np.uint8(np.random.randint(0,256,20000000))会临时创建比必要大4-8倍的数组。对于dtype=np.uint64，请使用np.int64(np.random.random_integers(0, _BLOCK_MAX, blocks_per_flush)).view(np.uint64)。 - jfs

@J.F.Sebastian：我找不到np.uint8()的文档，请问你有参考资料吗？我知道它可以将序列（也许是数组）转换为dtype uint8的数组，但是能否提供文档支持以证明这是一个稳定的特性呢？ - Eric O. Lebigot

@EOL: np.uint8在numpy中的行为与其他dtype类型相同。http://docs.scipy.org/doc/numpy/user/basics.types.html - jfs

Answer 2

由于您正在分配连续的块，因此应该能够执行以下操作（完全消除内部循环）：

for _ in xrange(0, num_flushes):
    in_memory_blocks[:blocks_per_flush] = numpy.random.randint(
            0, _BLOCK_MAX+1, blocks_per_flush)

    print('flushing bytes stored in memory...')

    # commented out for SO; exists in actual code
    # removing this doesn't make an order-of-magnitude difference in time
    # m.update(in_memory_blocks[:blocks_per_flush])

    in_memory_blocks[:blocks_per_flush].tofile(f)

这里使用了numpy.random.randint函数，它会分配一整块内存并填充随机整数（请注意J.F.Sebastian在下面的评论中提到的numpy.random.randint与random.randint的区别）。据我所知，没有办法使用numpy的随机程序来填充预先分配的数组。另一个问题是，numpy的randint返回int64数组。如果需要其他大小的整数，则可以使用numpy的类型方法，例如numpy.uint8。如果要生成覆盖类型的所有范围的randints，则@J. F. Sebastian使用numpy.random.bytes的方法将是最佳选择（几乎在任何情况下！）。

然而，简单的测试显示出合理的时间（与C代码的数量级相同）。以下代码测试使用numpy方法分配包含2000万个随机整数的uint8数组的时间：

from timeit import Timer
t = Timer(stmt='a=numpy.uint8(numpy.random.randint(0, 100, 20000000))',
        setup='import numpy')
test_runs = 50
time = t.timeit(test_runs)/test_runs
print time

我明白，在我的四年老的Core2笔记本电脑上，每个分配大约需要0.7秒（它运行50次，因此整个测试需要更长时间）。这是对20,000,000个随机uint8整数的每个分配所需的0.7秒，因此我预计整个500MB需要大约20秒。

更多的内存意味着您可以一次性分配更大的块，但是您仍然在有效地浪费时间为每个int分配和写入64位，而您只需要8位（我没有量化这种效应）。如果速度仍然不够快，您可以使用numpy ctypes接口调用C实现。这真的很容易使用，并且您几乎不会遇到纯C的减速。

总的来说，使用numpy时，始终尝试使用numpy例程，记住回退到ctypes的C也不太痛苦。总的来说，这种方法允许非常有效地使用Python进行数值处理，而且减速非常小。

编辑: 我想到了另一件事情: 如上所述的实现方式，我认为您将会多做一份不必要的拷贝。如果in_memory_blocks的长度为blocks_per_flush，那么最好直接将其分配给numpy.random.randint的返回值，而不是分配给某个子数组（在一般情况下必须是一个拷贝）。所以：

in_memory_blocks = numpy.random.randint(0, _BLOCK_MAX+1, blocks_per_flush)

而不是：

in_memory_blocks[:blocks_per_flush] = numpy.random.randint(
        0, _BLOCK_MAX+1, blocks_per_flush)

然而，经过计时后，第一种情况并没有显著提高速度（仅约2％），因此可能不值得过多担心。我猜绝大部分时间都花在了生成随机数上（这是我所期望的）。

- Henry Gomersall

---

这样做导致了惊人的加速。之前我只能达到每秒约1MB的速度。通过这个变化，我现在可以生成大约258MB / sec的速度。我注意到了numpy.random.randint函数，但跳过它是因为我错误地认为它是主要影响性能的内存分配。 - Allen George

numpy.random.randint()（半开区间）与 random.randint()（闭区间：两端都包含）不同。请使用 numpy.random.random_integers()。 - jfs

numpy.dtype('u8') 是 numpy.dtype('uint64')；它不是 `numpy.dtype('uint8')'。 - jfs

我的错误，我以为它们是一样的，我会添加澄清。 - Henry Gomersall

我已经在我的实现中删除了副本（即当我编写代码时，我使用了in_memory_blocks = numpy.random.randint(0, _BLOCK_MAX, blocks_per_flush)）。由于numpy.random.randint（...）返回一个新的缓冲区，所以我没有理由将其复制到in_memory_blocks。 - Allen George

Answer 3

如果你只是想每次以 block_size 字节的大小填充一个文件，这可能比之前的答案更快。基于生成器并完全绕过数组创建：

import numpy as np

def random_block_generator(block_size):
    while True:
        yield np.random.bytes(block_size)

rbg = random_block_generator(BLOCK_SIZE)

那么你的使用方法是：

f = open('testfile.bin','wb')

for _ in xrange(blocks_to_write):
    f.write( rbg.next())

f.close()

Numpy使用确定性随机数生成（序列中的下一个数字始终相同，只是在初始化时以随机位置开始）。如果您需要真正的随机数据（加密级别），则可以使用import Crypto.Random as cr和yield cr.get_random_bytes(block_size)而不是np。

此外，如果您的BLOCK_SIZE是一个定义好的常量，您可以使用类似于以下的生成器表达式（这次使用Crypto库）：

import Crypto.Random as cr
from itertools import repeat

BLOCK_SIZE = 1000

rbg = (cr.get_random_bytes(BLOCK_SIZE) for _ in repeat(0))

f = open('testfile.bin','wb')

for _ in xrange(blocks_to_write):
    f.write( rbg.next())

f.close()

这包括实现rbg=...和执行。即使带有稍微慢一些的Crypto.Random，这种生成器方法也会在达到计算极限之前耗尽磁盘I/O资源（虽然我相信其他答案也是如此）。

更新：

在 Athlon X2 245 上进行了一些计时数据--

• Crypto: 生成 500MB，不写入--10.8秒（46 MB/s）
• Crypto: 生成 500MB 并写入--11.2秒（44.5 MB/s）
• Numpy: 生成 500MB，不写入--1.4秒（360 MB/s）
• Numpy: 生成 500MB 并写入--7.1秒（70 MB/s）

因此，numpy 版本大约快 8 倍 （足以让我的旧盘驱动器达到极限）。我使用的是生成器表达式形式而不是生成器函数形式测试它们。

- dhj

Answer 4

我不太擅长优化，但我看不出你的代码能运行得更快。你使用了纯迭代器和O(1)访问结构。

问题在于你选择的编程语言。请记住，你正在虚拟机中运行，并且是解释器。你的C程序将始终运行快一个数量级。

- salezica

---

在Python中使用NumPy包实际上是很快的：耗时操作是通过编译代码完成的。因此，使用NumPy的C程序和Python程序可以具有类似的运行时间。虽然Python版本较慢（有时并不多），并且如果不使用快速的NumPy数组操作，则可能会慢得多。另一方面，我曾经看到Python程序比Fortran程序运行得更快，因为Fortran程序不必要地预先分配了巨大的数组，“以防万一”有大量数据需要处理。 - Eric O. Lebigot