如何加快python中的numpy数组填充？

Question

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

# 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 秒。相比之下，我使用 rand() 的简单 C 程序可以在 8 秒内创建一个 500MB 的文件。

如何提高上述循环的性能？我很确定我忽略了导致运行时巨大差异的明显因素。

Answer 1

由于您分配的是连续块，您应该能够执行以下操作（完全摆脱内部循环）：

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 涵盖类型的整个范围，那么下面使用 numpy.random.bytes 的@J. F. Sebastian 的方法将是最好的（几乎在任何情况下！）。

但是，简单的测试显示合理的时间（与 C 代码的数量级相同）。以下代码测试使用 numpy 方法分配 20,000,000 个随机整数的 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

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

更多内存意味着您可以一次分配更大的块，但是当您只需要 8 位时，您仍然有效地浪费时间为每个 int 分配和写入 64 位（我没有量化这种效果）。如果它仍然不够快，您可以使用 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%），因此可能不值得过多担心。我猜大部分时间都花在实际生成随机数上（这是我所期望的）。

Answer 2

由于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's one快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

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

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 上的一些计时数据 --

• 加密：生成 500MB，不要写入 -- 10.8s (46 MB/s)
• 加密：生成 500MB 并写入 -- 11.2s (44.5 MB/s)
• Numpy：生成 500MB，不要写入 -- 1.4s (360 MB/s)
• Numpy：生成 500MB，然后写入 -- 7.1s (70 MB/s)

所以 numpy 版本快了大约 8 倍（快到足以让我的旧盘片驱动器最大化）。我使用生成器表达式形式而不是生成器函数形式对它们进行了测试。

Answer 4

我不太擅长优化，但我看不出有什么方法可以让您的代码运行得更快。您正在使用纯迭代器和 O(1) 访问结构。

我认为，问题在于您选择的语言。请记住，您是在虚拟机中运行的，并且是一个解释器。您的 C 程序的运行速度总是会快一个数量级。