Batcher

class torchdata.datapipes.iter.Batcher(datapipe: IterDataPipe, batch_size: int, drop_last: bool = False, wrapper_class=List)

Creates mini-batches of data (functional name: batch). An outer dimension will be added as batch_size if drop_last is set to True, or length % batch_size for the last batch if drop_last is set to False.

Parameters:

• datapipe – Iterable DataPipe being batched

• batch_size – The size of each batch

• drop_last – Option to drop the last batch if it’s not full

• wrapper_class – wrapper to apply onto each batch (type List) before yielding, defaults to DataChunk

Example

>>> # xdoctest: +SKIP
>>> from torchdata.datapipes.iter import IterableWrapper
>>> dp = IterableWrapper(range(10))
>>> dp = dp.batch(batch_size=3, drop_last=True)
>>> list(dp)
[[0, 1, 2], [3, 4, 5], [6, 7, 8]]