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

Question

我正在尝试使用以下代码填充预分配的bytearray:

# 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需要~1s; 500MB需要~376s.相比之下,我使用rand()的简单C程序可以在8s内创建一个500MB的文件.

如何改善上述循环的性能？我很确定我忽略了一些显而易见的事情,这会导致运行时间的巨大差异.

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()如果数组的dtype已经存在,则不会复制np.int64..view(numpy.uint64)强制将其解释为无符号(也不执行复制).

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分配整个内存块并用随机整数填充它的函数（请注意下面 JF Sebastian 关于numpy.random.randintvs 的评论random.randint）。没有办法（据我所知）使用 numpy 随机例程填充预分配的数组。另一个问题是 numpy 的 randint 返回 int64 数组。如果您需要其他大小的整数，则可以使用 numpy 类型方法，例如 numpy.uint8。如果您希望 randints 覆盖类型的整个范围，那么下面使用 numpy.random.bytes 的@JF 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 时，您仍然有效地浪费时间为每个 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%），因此可能不值得过于担心。我想绝大多数时间都花在实际生成随机数上（这是我所期望的）。