python list 分批batch

python list 分批batch

def split_batch(init_list, batch_size):
    groups = zip(*(iter(init_list),) * batch_size)
    end_list = [list(i) for i in groups]
    count = len(init_list) % batch_size
    end_list.append(init_list[-count:]) if count != 0 else end_list
    return end_list


if __name__ == '__main__':

    a = [x for x in range(10)]
    init_batch_size = 5
    result = split_batch(a, init_batch_size)
    print(result)

结果

0-9的列表，按每5个分组
[[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]