基本上，我想根据给定列表 l 创建一个无限生成器批量大小为 batch_size .例如，如果我有 l = [1, 2, 3, 4, 5] 的列表和 batch_size = 2 , 我想生成 [1, 2] 的无限循环, [3, 4] , [5, 1] , [2, 3] , ...(类似于 itertool.circular 具有额外的批量大小)
我目前的方法如下，但还没有给出正确的解决方案，因为最后我只是在到达末尾时填充列表的第一个元素:
l = [1, 2, 3, 4, 5]
def generator(l, batch_size=2):
while 1:
for i in range(0, len(l), batch_size):
batch = l[i:(i + batch_size)]
if len(batch) < batch_size:
batch.extend(l[0: batch_size - len(batch)])
yield batch
>> gen = generator(l, batch_size=2)
>> next(gen) # [1, 2]
>> next(gen) # [3, 4]
>> next(gen) # [5, 1]
>> next(gen) # [1, 2] << want [2, 3] instead
有没有办法以循环的方式做到这一点？
请您参考如下方法:
是的，您基本上需要“take”和cycle的组合:
>>> def circle_batch(iterable, batchsize):
... it = itertools.cycle(iterable)
... while True:
... yield list(itertools.islice(it, batchsize))
...
>>> l = [1, 2, 3, 4, 5]
>>> c = circle_batch(l, 2)
>>> next(c)
[1, 2]
>>> next(c)
[3, 4]
>>> next(c)
[5, 1]
>>> next(c)
[2, 3]
>>> next(c)
[4, 5]
来自recipes in the docs您会看到 take 是一个基本工具，因此使用它:
>>> def take(n, iterable):
... "Return first n items of the iterable as a list"
... return list(islice(iterable, n))
...
>>> def cycle_batch(iterable, batchsize):
... it = itertools.cycle(iterable)
... while True:
... return take(batchsize, it)
...
>>> l = [1, 2, 3, 4, 5]
>>> c = circle_batch(l, 2)
>>> next(c)
[1, 2]
>>> next(c)
[3, 4]
>>> next(c)
[5, 1]
>>> next(c)
[2, 3]
>>> next(c)
[4, 5]
>>> next(c)
[1, 2]
