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 asbatch_size
ifdrop_last
is set toTrue
, orlength % batch_size
for the last batch ifdrop_last
is set toFalse
.- 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 toDataChunk
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]]