如何在PyTorch中将Python迭代器用作数据集？最优方案探讨

最直接高效的方案是用PyTorch内置的IterableDataset封装你的生成器，完全不需要自定义DataLoader，完美适配样本总量未知的流式场景。

核心思路 #

IterableDataset是PyTorch专门为流式、长度未知的数据集设计的接口，它不需要实现__len__方法，只需要实现__iter__来返回迭代器（也就是你的生成器），刚好匹配你的需求。

代码实现示例 #

修改你的原始代码，用IterableDataset包装生成器后再传入DataLoader：

import torch
from torch.utils.data import DataLoader, IterableDataset

def example_generator():
    for i in range(10):
        yield i

# 用IterableDataset封装生成器
class GeneratorDataset(IterableDataset):
    def __init__(self, gen_func):
        self.gen_func = gen_func

    def __iter__(self):
        # 返回你的生成器实例
        return self.gen_func()

BATCH_SIZE = 3
# 初始化数据集和DataLoader
train_dataset = GeneratorDataset(example_generator)
train_dataloader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=False)

# 直接迭代批量数据（不要调用len()，因为样本总量未知）
for batch_idx, batch in enumerate(train_dataloader):
    print(f"Batch {batch_idx}: {batch}")

关键注意事项 #

1. 不要调用len(train_dataloader)：因为IterableDataset没有总长度信息，强行调用会触发报错，直接迭代即可，这也符合你“样本总量未知”的需求。
2. 多进程加载的适配：如果需要设置num_workers>0，要避免多个worker重复生成数据。可以在__iter__里通过get_worker_info()获取worker ID，让每个worker生成不同分片的数据：

   def __iter__(self):
        worker_info = torch.utils.data.get_worker_info()
        if worker_info is None:
            # 单进程模式，直接返回完整生成器
            return self.gen_func()
        else:
            # 多进程模式，根据worker ID生成对应分片的数据
            # 这里需要你的生成器支持分片逻辑，比如传入start/offset参数
            return self.gen_func(start=worker_info.id, total_workers=worker_info.num_workers)
   

   如果你的生成器是完全流式、无法分片的，建议保持num_workers=0。

内容的提问来源于stack exchange，提问作者bja