PyTorch IterableDataset 怎样为每个 worker 分配不同数据片段

用 PyTorch 的 IterableDataset 配合 DataLoader(num_workers>0) 时，最容易误判的一点是：数据集对象会在每个 worker 进程里各有一份副本。如果 __iter__() 仍然从同一个起点完整遍历，每个 worker 都会读到同样的样本。

解决办法是把分片逻辑放进 __iter__()：通过 torch.utils.data.get_worker_info() 取得当前 worker 的编号和总数，再把整体的半开区间 [start, end) 切成互不重叠的局部区间。单进程运行时，get_worker_info() 返回 None，代码仍然保持原来的完整遍历。

核心判断：IterableDataset 不会自动替你做 worker 分片。只要一个数据集副本没有根据 worker id 改变读取范围，多 worker 就可能重复消费同一批数据。

这篇只讨论 DataLoader 的多进程 worker 分片，不展开多机或分布式训练的采样器选型。

为什么多 worker 会读到重复样本

对普通的 map-style 数据集，DataLoader 可以由主进程生成索引，再把索引分发给 worker。IterableDataset 没有这种随机索引模型，数据是由各个副本自行迭代产生的。于是设置 num_workers=2 后，Worker 0 和 Worker 1 都可能从 0 开始读取。

因此问题不在于 DataLoader “把同一条样本发了两次”，而在于两个独立的迭代器都认为自己拥有完整数据源。最小的修复边界是：每个 worker 只负责一个不重叠的区间或文件分片，并且这个判断要发生在 worker 进程执行 __iter__() 时。

在 __iter__ 中按 worker 切分区间

下面用整数区间模拟真实数据源。start 和 end 采用半开区间，意味着包含起点、不包含终点；这样相邻区间可以直接拼接，不会重复边界值。

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


class RangeDataset(IterableDataset):
    def __init__(self, start: int, end: int):
        super().__init__()
        self.start = start
        self.end = end

    def __iter__(self):
        # 主进程没有 worker 信息；多进程时每个 worker 都会拿到自己的副本。
        worker_info = get_worker_info()
        local_start, local_end = self.start, self.end

        if worker_info is not None:
            # 用向上取整给每个 worker 一个连续的半开区间。
            total = local_end - local_start
            per_worker = int(math.ceil(total / worker_info.num_workers))
            local_start = self.start + worker_info.id * per_worker
            local_end = min(local_start + per_worker, self.end)

        # 空区间自然不会产生样本，避免为了补齐数量而重复读取。
        yield from range(local_start, local_end)


dataset = RangeDataset(0, 10)
loader = DataLoader(dataset, num_workers=2, batch_size=2)
for batch in loader:
    # 这里关注的是不重复分片，DataLoader 返回顺序仍由调度和批次组成决定。
    print(batch)

以 [0, 10) 和两个 worker 为例，per_worker 为 5，两个副本分别负责 [0, 5) 与 [5, 10)。当总数不能被 worker 数整除时，向上取整会让前面的区间先拿到一个完整容量，最后一个区间可能更短，但不会越过 end。

不整除、尾部数据和 drop_last 怎么判断

分片正确只说明“同一条样本不会被两个 worker 主动读到”，不代表每个 batch 都完整。IterableDataset 在多进程读取时，每个 worker 都有自己的数据副本，drop_last=True 可能按副本丢弃各自最后一个不完整 batch。因此总样本数、实际 batch 数和最后一批大小不能只用全局总数简单相除。

排查时按下面三点看：

• 先把 num_workers 设为 0，确认单进程数据源本身没有重复。
• 再用一个明确的有限区间，记录每个 worker 的 id、局部起点和局部终点，确认区间相邻且不重叠。
• 最后再加入真实的 batch_size 和 drop_last，把“分片重复”和“尾批被丢弃”分开判断。

如果真实数据源不是整数范围，而是文件列表、数据库游标或远程对象，也要把“打开资源”和“选择片段”放在 worker 可见的边界内。不要在主进程创建一个不可安全复制的连接对象，再指望多个副本共享它完成分片。

上线前的四个检查点

第一，确认分片依据是 worker 的稳定身份，而不是在多个进程间共享并自增的普通 Python 变量。第二，使用 self.start、self.end 计算局部变量，避免一次数据集对象被重复创建迭代器时留下错误的上次状态。第三，明确尾部策略：允许短尾批次就保留，必须整批才处理就显式设置 drop_last 并接受其丢弃规则。第四，先用可数且容易核对的样本源做去重检查，再接入无限流或远程数据源。

如果改成 worker_init_fn 来配置每个副本，也要保持同一原则：每个副本必须获得不同的读取边界；只初始化随机种子而不改变数据范围，并不能解决重复消费。

常见问题

单进程时为什么不需要特殊分片？

因为 get_worker_info() 在主进程返回 None，示例会直接使用完整的 [start, end)。这也是先用 num_workers=0 排查数据源的原因。

worker 数量变化后能继续使用旧的区间吗？

不能把旧的静态边界硬编码下来。区间应当在每次 __iter__() 中根据当前的 num_workers 和 worker_id 重新计算，否则 worker 数变化后可能出现遗漏或重叠。

为什么输出顺序看起来不是按区间连续排列？

多个 worker 会并行产出，DataLoader 还会按 batch 组织结果。输出顺序不是判断分片是否重复的可靠依据；应检查每个 worker 的局部范围和全局样本集合。

记住一句话：IterableDataset 负责“怎样产生数据”，DataLoader worker 负责“在哪个进程产生数据”，两者之间的分片契约需要由 __iter__() 显式补上。