BatchMapper

class torchdata.datapipes.iter.BatchMapper(datapipe: IterDataPipe, fn: Callable, batch_size: int, input_col=None)

将源 DataPipe 中的元素组合成批次，并对每个批次应用函数，然后将输出展平为单个、未嵌套的 IterDataPipe（函数名称：map_batches）。

参数:

• datapipe – 源 IterDataPipe

• fn – 要应用于每个数据批次的函数

• batch_size – 从 datapipe 中聚合的批次大小

• input_col –

  应用 fn 的数据索引或索引，例如

  • None 作为默认值，将 fn 直接应用于数据。

  • 整数用于列表/元组。

  • 键用于字典。

示例

>>> from torchdata.datapipes.iter import IterableWrapper
>>> def fn(batch):
>>>     return [d + 1 for d in batch]
>>> source_dp = IterableWrapper(list(range(5)))
>>> mapped_dp = source_dp.map_batches(fn, batch_size=3)
>>> list(mapped_dp)
[1, 2, 3, 4, 5]

备注

与 map 相比，map_batches 不带 output_col 参数的原因是 fn 输出的大小不能保证与输入批次相同。由于大小不同，此操作无法将数据分配回原始数据结构。

并且，此操作是根据 TorchText 中的用例引入的。可以应用 pybinded C++ 矢量化函数以提高效率。