我想分块输入流以进行批处理。给定一个输入列表或生成器,
x_in = [1, 2, 3, 4, 5, 6 ...]
我想要一个可以返回该输入块的函数。说,如果chunk_size=4
,那么,
x_chunked = [[1, 2, 3, 4], [5, 6, ...], ...]
这是我一遍又一遍地做的事情,并且想知道是否有比自己编写更标准的方法。我错过了什么itertools
吗?enumerate
(可以用and解决这个问题groupby
,但感觉很笨拙。)如果有人想看到一个实现,这里是,
def chunk_input_stream(input_stream, chunk_size):
"""partition a generator in a streaming fashion"""
assert chunk_size >= 1
accumulator = []
for x in input_stream:
accumulator.append(x)
if len(accumulator) == chunk_size:
yield accumulator
accumulator = []
if accumulator:
yield accumulator
编辑
受 kreativitea 回答的启发,这里有一个解决方案islice
,它很简单且不需要后过滤,
from itertools import islice
def chunk_input_stream(input_stream, chunk_size):
while True:
chunk = list(islice(input_stream, chunk_size))
if chunk:
yield chunk
else:
return
# test it with list(chunk_input_stream(iter([1, 2, 3, 4]), 3))