batch.py 文件源码

def shuffle_and_batch(items: List[T], batch_size: int,
                      rng: Optional[random.Random] = None) \
        -> Iterator[List[T]]:
    """Optionally shuffles and batches items in a list.

    Args:
        - items: List of items to shuffle & batch.
        - batch_size: size of batches.
        - rng: random number generator if items should be shuffles, else None.

    Returns: Batch iterator
    """

    todo = list(range(len(items)))
    if rng is not None:
        rng.shuffle(todo)
    while todo:
        indices = todo[:batch_size]
        todo = todo[batch_size:]
        items_batch = [items[i] for i in indices]
        yield items_batch