多worker DataLoader中IterableDataset应报告何种大小？

为什么每个worker都会遍历完整数据集？ #

IterableDataset的默认逻辑是每个worker会独立初始化并遍历整个数据集。它和MapDataset不同：MapDataset基于索引访问数据，PyTorch会自动给每个worker分配不同的索引区间，避免重复；但IterableDataset是按迭代器方式逐个返回数据，PyTorch无法自动对其进行分片，所以默认情况下每个worker都会从头遍历到尾，最终总样本数等于原数据集大小乘以worker数量。

如何在各worker中配置数据集？ #

正确做法是给每个worker分配原数据集的一个分片，需要利用torch.utils.data.get_worker_info()获取当前worker的ID和总worker数，然后在迭代时只返回对应分片的数据。

修改后的示例代码：

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


class Dataset(IterableDataset):
    def __init__(self, size: int):
        super().__init__()
        self.size = size

    def __iter__(self):
        worker_info = torch.utils.data.get_worker_info()
        if worker_info is None:
            # 单worker模式，遍历全部数据
            iter_start = 0
            iter_end = self.size
        else:
            # 多worker模式，按worker ID拆分数据分片
            per_worker = self.size // worker_info.num_workers
            iter_start = worker_info.id * per_worker
            # 处理不能整除的情况，最后一个worker承担剩余样本
            iter_end = iter_start + per_worker if worker_info.id != worker_info.num_workers - 1 else self.size
        for i in range(iter_start, iter_end):
            yield torch.tensor([i], dtype=torch.float32)

    def __len__(self):
        return self.size


dataloader = DataLoader(Dataset(5), batch_size=1, num_workers=4)
result = list(dataloader)
print(len(result))  # 输出5，符合原数据集大小

应该修改数据集长度还是拆分每个worker的数据集大小？ #

两种方式都不对：

• 不要把__len__改成self.size*4，这会让DataLoader错误预期20个样本，而实际分片后只有5个，依然会触发长度不匹配警告。
• 正确做法是保持__len__返回原数据集的真实总大小，同时让每个worker只遍历自己负责的分片（原大小的1/4左右，处理整除边界）。这样DataLoader的预期样本数和实际产出一致，不会有警告。

为什么DataLoader会出现长度不匹配的情况？ #

不是DataLoader期望len(dataset)等于样本数×worker数，而是默认情况下每个worker都会遍历完整的IterableDataset，导致实际产出样本数是len(dataset)*num_workers，但DataLoader会用dataset.__len__()的值作为预期总样本数，两者不匹配就会抛出警告。MapDataset不存在这个问题，因为PyTorch会自动处理索引分片，而IterableDataset需要用户手动实现分片逻辑，否则就会重复遍历。

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