在numpy数组中循环使用python生成器

Question

我的代码中的速度瓶颈是两个数组x和y的元素循环的两倍.用于提高性能的标准hpc技巧是以块的形式进行循环,以便可以最小化缓存未命中.我正在尝试使用python生成器来进行分块,但是需要在外部for循环中不断重新创建耗尽的生成器,这会杀死我的运行时.

题:

是否有更智能的算法来构造适当的生成器来执行chunked double-for循环？

具体插图:

我将创建两个虚拟数组,x和y.我会简短地说明它们,但实际上这些是带有~1e6元素的numpy数组.

x = np.array(['a', 'b', 'b', 'c', 'c', 'd'])
y = np.array(['e', 'f', 'f', 'g'])

天真的双循环只是:

for xletter in x:
    for yletter in y:
        # algebraic manipulations on x & y

现在让我们使用生成器以块的形式执行此循环:

chunk_size = 3
xchunk_gen = (x[i: i+chunk_size] for i in range(0, len(x), chunk_size))
for xchunk in xchunk_gen:
    ychunk_gen = (y[i: i+chunk_size] for i in range(0, len(y), chunk_size))
    for ychunk in ychunk_gen:
        for xletter in xchunk:
            for yletter in ychunk:
                # algebraic manipulations on x & y

请注意,为了实现此问题的生成器解决方案,我必须在外部循环中不断重新创建ychunk_gen.由于y是一个大型数组,这会破坏我的运行时间(对于~1e6元素,在我的笔记本电脑上创建这个生成器大约需要20ms).

有没有办法更巧妙地说明我如何构建解决这个问题的生成器？或者是否有必要完全抛弃发电机解决方案？

(注意:在实践中,我使用cython来执行这个紧密循环,但无论如何都适用于上述所有情况).

Answer 1

在我看来，生成器表达式的创建会浪费您的运行时间，因为它没有被 cython 优化。

一个更好的解决方案是使用numexpr来处理所有缓存优化问题。由于 x 和 y 的操作是代数，因此它应该非常适合您的约束（numexpr 可以做得更多）