pytorch 中对 iterabledataset 使用 num_workers > 0 时，默认每个工作进程独立执行完整迭代，导致数据重复；需通过 get_worker_info() 手动划分数据范围，确保各进程只处理子集。
pytorch 中对 iterabledataset 使用 num_workers > 0 时，默认每个工作进程独立执行完整迭代，导致数据重复；需通过 get_worker_info() 手动划分数据范围，确保各进程只处理子集。
在 PyTorch 中，IterableDataset 的设计初衷是支持流式、无限或无法预知长度的数据源（如日志流、数据库游标、远程文件读取等），因此它不依赖 __len__ 进行索引切分，而是完全由 __iter__() 方法定义数据生成逻辑。正因如此，当配合 DataLoader 并启用多进程（num_workers > 0）时，每个 worker 会实例化一份独立的 dataset 对象，并各自调用其 __iter__() —— 若未做任何区分，所有 worker 将产出完全相同的数据序列，造成批次重复、数据量翻倍甚至更严重的问题。
例如原始代码中：
class ToyDataset(torch.utils.data.IterableDataset):
def __iter__(self):
data = torch.arange(len(self)) # 每个 worker 都生成 0~385 全量
yield from data
def __len__(self): return 386
启用 num_workers=2 后，两个 worker 分别产出全部 386 个样本，经 batch_size=256 后得到 ceil(386/256)=2 批 × 2 个 worker = 实际迭代出 4 个 batch（len(list(loader)) == 4），而预期应为 2 批。
✅ 正确做法是：利用 torch.utils.data.get_worker_info() 在 __iter__() 中识别当前 worker 身份，并按 worker ID 划分全局数据区间，实现数据分片（sharding）。关键要点如下：
-
get_worker_info()在主进程返回None，在 worker 进程中返回包含id（从 0 开始）、num_workers等信息的WorkerInfo对象； - 数据划分应满足：互斥 + 并集覆盖全集 + 均衡（尽可能）；
- 推荐使用
math.ceil((end - start) / num_workers)计算每份大小，避免整除误差导致数据丢失； - 初始数据范围建议显式传入
__init__（如start=0, end=386），而非依赖__len__，增强可控性与语义清晰度。
以下是修正后的完整示例：
import torch
import math
class ShardedToyDataset(torch.utils.data.IterableDataset):
def __init__(self, start: int, end: int):
super().__init__()
self.start = start
self.end = end
def __iter__(self):
worker_info = torch.utils.data.get_worker_info()
if worker_info is None:
# 主进程（单进程模式）或未启用多进程时，使用全量
iter_start, iter_end = self.start, self.end
else:
# 多进程下：按 worker_id 切分 [start, end)
per_worker = int(math.ceil((self.end - self.start) / worker_info.num_workers))
worker_id = worker_info.id
iter_start = self.start + worker_id * per_worker
iter_end = min(iter_start + per_worker, self.end)
# 生成本 worker 负责的子序列
for x in range(iter_start, iter_end):
yield torch.tensor(x, dtype=torch.long)
# 使用示例
dataset = ShardedToyDataset(0, 386)
loader = torch.utils.data.DataLoader(
dataset,
batch_size=256,
num_workers=2,
persistent_workers=True # 可选：避免 worker 重复启停开销
)
batches = list(loader)
print(f"Total batches: {len(batches)}") # 输出: Total batches: 2
print(f"Batch shapes: {[b.shape for b in batches]}") # e.g., [torch.Size([256]), torch.Size([130])]
⚠️ 注意事项：
-
勿在
__iter__()中使用随机种子重置（除非明确需要每个 worker 独立随机性），否则可能破坏可复现性； - 若数据源本身不可分片（如单个实时 API 流），应改用
num_workers=0或自定义worker_init_fn实现 token/offset 同步； -
persistent_workers=True（PyTorch ≥ 1.7）可减少 worker 进程反复创建开销，但需确保 dataset 实例线程安全； -
IterableDataset不支持shuffle=True（会报错），如需打乱，应在__iter__()内部对本地分片做 shuffle（注意：全局 shuffle 无法保证）。
总结：IterableDataset + 多进程不是“开箱即用”的组合，而是需要开发者主动承担数据分片责任。理解 get_worker_info() 的作用机制并合理划分数据边界，是构建高效、正确分布式数据加载 pipeline 的关键一步。
