批处理生成器函数 returns 转换为列表时仅最后一个元素

Batch Generator function returns only last element when converted to list

我有以下代码从项目列表中生成批次:

def batch_generator(items, batch_size):
    count = 1
    chunk = []
    
    for item in items:
        if count % batch_size:
            chunk.append(item)
        else:
            chunk.append(item)
            yield chunk
            chunk.clear()
        count += 1
    
    if len(chunk):
        yield chunk

逐个迭代产生预期结果:

for x in batch_generator(range(17), 5):
    print(x)
# [0, 1, 2, 3, 4]
# [5, 6, 7, 8, 9]
# [10, 11, 12, 13, 14]
# [15, 16]

然而,当我将生成器直接转换为列表时,多次只返回最后一个元素!

list(batch_generator(range(17), 5))
# [[15, 16], [15, 16], [15, 16], [15, 16]]

而转换为列表的简单生成器工作正常:

list(([i,i*2,i*3] for i in range(5)))
# [[0, 0, 0], [1, 2, 3], [2, 4, 6], [3, 6, 9], [4, 8, 12]]

为什么会这样?

chunk.clear() 是这里的问题。归根结底,返回的列表是多次返回的同一个列表。

chunk.clear() 替换为 chunk = []。这样 chunk 将是列表的不同实例:

您可以使用 chunk[:]list(chunk) 生成块的副本。而不仅仅是 yield chunk.