torchdata DataPipe 和 DataLoader2 快速入门和使用

torchdata 使用了 dill 替代了 pickle， 从而允许在spawn多进程模式下转移复杂的数据集。

是该博文的另一种解决方法
pytorch 使用 multiprocess库 让 Dataloader 不再报错 AttributeError: Can‘t pickle local object
https://blog.csdn.net/ONE_SIX_MIX/article/details/130388252

以下为使用 torchdata 的示例，SimpleDataset 可以直接替换为你自己的Dataset

from torchdata.datapipes.iter import IterableWrapper
from torchdata.dataloader2 import DataLoader2
from torchdata.dataloader2 import MultiProcessingReadingService


# 定义一个简单数据集，包含lambda函数
class SimpleDataset:
    def __init__(self):
        # lambda function here
        self.func = lambda x: x+1

    def __len__(self):
        # 必须实现
        return 1000

    def __getitem__(self, i):
        # 必须实现
        if i >= len(self):
            raise StopIteration()
        return self.func(i)


def map_fn(x):
    # 转换为 torch.Tensor
    # 要在里面 import torch 才能使用 torch 的函数，不知为何外面的 torch 导入被忽略了。
    import torch
    return torch.tensor(x)


if __name__ == '__main__':

    batchsize = 2
    n_worker = 4

    # 定义一个数据集
    ds = SimpleDataset()

    # 使用序列包装后，才能给 DataLoader2 使用
    warp_ds = IterableWrapper(ds)

    # 使用多进程读取后端
    rs = MultiProcessingReadingService(num_workers=n_worker)

    # 数据加载器。对数据集使用 混洗shuffle->对每个样本进行map_fn操作map->批量化batch->按pytorch的dalaloader的默认方式进行分组collate->使每个进程生成的样本唯一sharding_filter
    dl = DataLoader2(warp_ds.shuffle().map(map_fn).batch(batchsize).collate().sharding_filter(), reading_service=rs)

    for x in dl:
        print(x)