gen_schedule/gen_utils.py

import numpy as np

def even_split(origin_list, batch_size):
    '''
    uniformly split a list into little batches
    :param origin_list:
    :param batch_size:
    :return:
    '''
    total_size = len(origin_list)
    n_batch = np.ceil(total_size / batch_size).astype(np.int32)
    avg_size = total_size // n_batch
    target_size = [int(avg_size + 1 if i < total_size % n_batch else avg_size) for i in range(n_batch)]
    splatted = []
    cnt = 0
    for t_size in target_size:
        splatted.append(origin_list[cnt: cnt + t_size])
        cnt += t_size
    return splatted