Python生成器在紧凑的双重for循环中有效地使用于numpy数组。

Question

我的代码中速度瓶颈是在一个双重循环中，遍历两个数组x和y的元素。提高性能的标准HPC技巧是对循环进行分块，以最小化缓存未命中。我试图使用Python生成器来进行分块，但需要在外层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个元素，在我的笔记本电脑上创建此生成器需要约20毫秒）。

有没有更聪明的方式来构建生成器，以避免这个问题？或者放弃使用生成器方案是必要的吗？

（注：在实践中，我正在使用Cython执行这个紧密的循环，但以上所有内容都适用）。

- aph

---

如果x和y是RAM中的列表，那么像你这样使用生成器并没有任何好处......另外，你可以直接运行counter = len(x) * len(y)。 - Jörn Hees

此外，您能告诉我们您在实际循环中正在做什么以及您的实际任务是什么吗？也许如果我们了解整体情况，就可以提供更好的帮助。 - Stefan Pochmann

@PM2Ring - 是的，实际上我正在使用连续的numpy数组，正是出于这个原因。 - aph

顺便说一句，我并不是在说你的问题有问题，事实上它本身就很好而且很有趣，回答这个问题可以帮助你。我只是想说，你直接询问你实际遇到的问题可能会更好。 - Stefan Pochmann

你不能直接使用C风格迭代来模拟分块生成器吗？ - hpaulj

Answer 1

在我看来，您的生成器表达式的创建正在消耗您的运行时间，因为它没有通过cython进行优化。

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

- Xavier Combelle

Answer 2

你正在xchunk循环中重新定义ychunk_gen。也许以下方法会更快：

chunk_size = 3
xchunk_gen = (x[i: i+chunk_size] for i in xrange(0, len(x), chunk_size))

def ychunk_gen(some_dependency_on_outer_loop):
    # use some_dependency_on_outer_loop
    for i in xrange(0, len(y), chunk_size):
        yield y[i: i+chunk_size]

for xchunk in xchunk_gen:
    for ychunk in ychunk_gen(chunk_or_something_else):
        for xletter in xchunk:
            for yletter in ychunk:
                # algebraic manipulations on x & y

也许有一种更好的方法：

我假设 x 和 y 是 numpy 数组，因此您可以重塑数组，然后循环遍历每一行：

for xchunk in x.reshape((len(x)//chunk_size, chunk_size)):
    for ychunk in y.reshape((len(y)//chunk_size, chunk_size)):
        # the letter loops

我在http://docs.scipy.org/doc/numpy/reference/generated/numpy.reshape.html中读到，如果你不想通过reshape复制数据，你应该改变数据的shape属性：

x.shape = len(x)//chunk_size, chunk_size 
y.shape = len(y)//chunk_size, chunk_size

- koffein

---

这个第二个numpy解决方案如果chunk_size不能整除数组长度，是否有效？看起来似乎不行，因为reshape对于数组长度的保留有严格要求，但也许有些我不理解的东西？ - aph

你对形状限制是正确的，但我认为这很容易解决？ - koffein

Answer 3

itertools.tee 可以节省一些时间：

y = np.arange(100)
def foo1(y):
   # create ygen each loop
   # py3 so range is a generator
   for j in range(100):
       ygen=(y[i:i+10] for i in range(0,1000,10))
       r = [x.sum() for x in ygen]
   return r

def foo3(y):
   # use tee to replicate the gen
   ygen=(y[i:i+10] for i in range(0,1000,10))
   ygens=itertools.tee(ygen,100)
   for g in ygens:
       r=[x.sum() for x in g]
   return r

In [1123]: timeit foo3(y)
10 loops, best of 3: 108 ms per loop
In [1125]: timeit foo1(y)
10 loops, best of 3: 144 ms per loop

但基于

http://docs.cython.org/0.15/src/userguide/limitations.html#generators-and-generator-expressions

自从Cython 0.13版本以来，一些生成器表达式在与内置函数结合时可以转换为内联循环的形式得到支持，例如sum(x*2 for x in seq)。截至0.14版本，支持的内置函数有list()、set()、dict()、sum()、any()、all()和sorted()。

我想知道cython对你的分块生成器表达式做了什么。

---

重新塑形和迭代行并不能很好地节省时间。

def foo4(y):
   y2d=y.reshape(100,10)
   for _ in range(100):
       r=[x.sum() for x in y2d]
   return r

相对于 teed 生成器，速度稍慢一些。当然，这样的相对时间可能会随着数组大小的变化而改变。

- hpaulj