Python itertools - 慢吗？

Question

我正在尝试使用Python的itertools模块来加速三重嵌套的for循环。下面的测试代码比较了标准的三重嵌套循环和itertools的product方法，并输出：

嵌套循环时间 = 2.35023秒

Itertools循环时间 = 2.67766秒

我是否漏掉了什么？

import numpy
import itertools
import time

n = 128
a = numpy.arange(n**3).reshape((n,n,n))
b = numpy.zeros((n,n,n))
c = numpy.zeros((n,n,n))

t = time.time()
for i in range(n):
    for j in range(n):
        for k in range(n):
            b[i,j,k] = a[i,j,k]
print 'Nested loop time = %g secs' % (time.time() - t)

t = time.time()
for (i,j,k) in itertools.product(range(n), repeat=3):
    c[i,j,k] = a[i,j,k]
print 'Itertools loop time = %g secs' % (time.time() - t)

- repoman

---

我有什么遗漏吗？-- 你似乎忽略了一个事实，即没有人声称 itertools.product() 的想法是为了加速嵌套的for循环。 - Sven Marnach

9.7 itertools-创建迭代器以实现高效循环的函数...http://docs.python.org/library/itertools.html - Robert William Hanks

Answer 1

似乎在较大的 n 值下，itertools.product 速度较慢：

In [24]: print _23
from itertools import product
def nested_loops(n):
    for i in range(n):
        for j in range(n):
            for k in range(n):
                pass
def itertools_product(n):
    for (i,j,k) in product(range(n), repeat=3):
        pass
In [25]: %timeit nested_loops(128)
10 loops, best of 3: 68.6 ms per loop

In [26]: %timeit itertools_product(128)
10 loops, best of 3: 162 ms per loop

In [27]: %timeit nested_loops(10)
10000 loops, best of 3: 84.5 us per loop

In [28]: %timeit itertools_product(10)
10000 loops, best of 3: 79.8 us per loop

In [30]: %timeit nested_loops(300)
1 loops, best of 3: 833 ms per loop

In [31]: %timeit itertools_product(300)
1 loops, best of 3: 2.07 s per loop

不使用元组拆包：

In [40]: print _39
from itertools import product
def itertools_product(n):
    for ijk in product(range(n), repeat=3):
        pass
In [41]: %timeit itertools_product(128)
10 loops, best of 3: 115 ms per loop

In [42]: %timeit itertools_product(10)
10000 loops, best of 3: 59.2 us per loop

In [43]: %timeit itertools_product(300)
1 loops, best of 3: 1.47 s per loop

此外，出于好奇，使用列表推导和生成器表达式：

def list_comprehension_product(n):
    range_n = range(n)
    for (i,j,k) in [ (i, j, k) for i in range_n for j in range_n for k in range_n ]:
        pass
def generator_expression_product(n):
    range_n = range(n)
    for (i,j,k) in ( (i, j, k) for i in range_n for j in range_n for k in range_n ):
        pass
In [51]: %timeit list_comprehension_product(128)
1 loops, best of 3: 583 ms per loop

In [52]: %timeit generator_expression_product(128)
1 loops, best of 3: 480 ms per loop

这些基准测试是在 python --version 下运行的：

2.6.7 (r267:88850, Jul 31 2011, 19:30:54) 
[GCC 4.2.1 (Based on Apple Inc. build 5658) (LLVM build 2335.15.00)]

- David Wolever

---

是的，对于大范围来说似乎会慢一些。而这正是你希望它更快的地方 - 对吧？ - repoman

Answer 2

看起来第二个循环比第一个慢，可能是因为元组解包的缘故。你不必这样做，我发现像这样做可以使第二个循环更快：

for ijk in itertools.product(range(n), repeat=3):
    c[ijk] = a[ijk]

当然，使用numpy时，您应该尽量避免循环遍历元素，而是一次性在整个数组上使用numpy操作。这样，所有的循环等操作都是在C语言中完成的，您将获得巨大的速度提升。

- Ned Batchelder

---

不想让这与numpy有关。只是使用numpy来快速进行三重索引。本可以调用某个带有三重参数的函数... - repoman