Python "restart"生成器与 itertools.cycle 或类似构造一起使用时？

Question

我想从生成器生成随机值，并且希望能够无限次生成它们，因此我将生成器包装在 itertools.cycle 中。

下面是一个示例实现。

from itertools import cycle
import numpy as np

def generator(batch_size):
  for _ in range(batch_size):
    yield np.random.randint(0,50)

np.random.seed(111)
batch_size = 3
gen = cycle(generator(batch_size))

for _ in range(10):
  print([next(gen) for _ in range(batch_size)])

输出:

[20, 44, 20]
[20, 44, 20]
[20, 44, 20]
[20, 44, 20]
[20, 44, 20]
[20, 44, 20]
[20, 44, 20]
[20, 44, 20]
[20, 44, 20]
[20, 44, 20]

请注意，输出只是重复相同的三个数字。

我的直觉是发电机应该在每个周期开始时“重新启动”，但事实似乎并非如此。相反，生成器似乎运行一次以生成初始列表，然后无限期地重复输出。

实现所需行为的另一种策略是放弃cycle，而是将infinite=True参数传递给生成器，然后使用while True 循环而不是 for _ in range(...) 循环:

def generator(batch_size, infinite):
  while True:
    for _ in range(batch_size):
      yield np.random.randint(0,50)
    if not infinite:
      break

但这感觉很笨拙。

是否有任何构造可以用来包装生成器，以便它实际上在每个周期都重新启动？

Answer 1

构建并扁平化无限的生成器流:

from itertools import chain, repeat
infinite_gen = chain.from_iterable(map(generator, repeat(batch_size)))
#                 flatten            infinite stream of generators

或者在 Python 2 上:

from itertools import chain, repeat, imap
infinite_gen = chain.from_iterable(imap(generator, repeat(batch_size)))