Python PyTorch ParallelReadConcat用法及代碼示例

本文簡要介紹python語言中 torchrec.datasets.utils.ParallelReadConcat 的用法。

用法:

class torchrec.datasets.utils.ParallelReadConcat(*datapipes: torch.utils.data.dataset.IterDataPipe, dp_selector: typing.Callable[[typing.Sequence[torch.utils.data.dataset.IterDataPipe]], typing.Sequence[torch.utils.data.dataset.IterDataPipe]] = <function _default_dp_selector>)

基礎：torch.utils.data.dataset.IterDataPipe

ParallelReadConcat 。

可迭代 DataPipe 連接多個可迭代 DataPipe。與 DataLoader 一起使用時，將數據管道的子集分配給每個 DataLoader 工作線程以允許並行讀取。 :param datapipes: IterDataPipe 要讀取的實例。 :param dp_selector：每個DataLoader工作人員將使用該函數來確定要讀取的數據管道的子集：param：

例子：

datapipes = [
    criteo_terabyte(
        (f"/home/local/datasets/criteo/shard_{idx}.tsv",),
    )
    .batch(100)
    .collate()
    for idx in range(4)
]
dataloader = DataLoader(
    ParallelReadConcat(*datapipes), num_workers=4, batch_size=None
)