Python PyTorch Collator用法及代码示例

本文简要介绍python语言中 torchdata.datapipes.iter.Collator 的用法。

用法:

class torchdata.datapipes.iter.Collator(datapipe: IterDataPipe, collate_fn: Callable = <function default_collate>)

参数：

• datapipe-正在整理可迭代DataPipe

• collate_fn-自定义整理函数，收集和合并数据或一批数据。默认函数根据数据类型整理到张量。

通过自定义整理函数(函数名称：collate)将样本从 DataPipe 整理到张量。默认情况下，它使用 torch.utils.data.default_collate() 。

注意

在编写自定义整理函数时，您可以导入 torch.utils.data.default_collate() 以获取默认行为，并导入 functools.partial 以指定任何其他参数。

示例：将整数数据转换为浮点张量

>>> class MyIterDataPipe(torch.utils.data.IterDataPipe):
...     def __init__(self, start, end):
...         super(MyIterDataPipe).__init__()
...         assert end > start, "this example code only works with end >= start"
...         self.start = start
...         self.end = end
...
...     def __iter__(self):
...         return iter(range(self.start, self.end))
...
...     def __len__(self):
...         return self.end - self.start
...
>>> ds = MyIterDataPipe(start=3, end=7)
>>> print(list(ds))
[3, 4, 5, 6]
>>> def collate_fn(batch):
...     return torch.tensor(batch, dtype=torch.float)
...
>>> collated_ds = CollateIterDataPipe(ds, collate_fn=collate_fn)
>>> print(list(collated_ds))
[tensor(3.), tensor(4.), tensor(5.), tensor(6.)]