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.)]