pytorch - 如何将生成器转换为 Pytorch 数据加载器？

Question

我有一个生成合成数据的生成器。如何将其转换为 PyTorch 数据加载器？

Answer 1

你可以用 data.IterableDataset 包装你的生成器:

class IterDataset(data.IterableDataset):
    def __init__(self, generator):
        self.generator = generator

    def __iter__(self):
        return self.generator()

自然地，您可以用 data.DataLoader 包装这个数据集.

这是一个显示其用法的最小示例:

>>> gen = lambda: [(yield x) for x in range(10)]

>>> dataset = IterDataset(gen)
>>> for i in data.DataLoader(dataset, batch_size=2):
...    print(i)
tensor([0, 1])
tensor([2, 3])
tensor([4, 5])
tensor([6, 7])
tensor([8, 9])