Dataset 与 DataLoader

torch.utils.data 用两套协议把「数据从哪来」与「数据怎么喂」解耦：Dataset 描述单样本读取，DataLoader 负责批组装、乱序、多进程预取与搬运，是 PyTorch 架构总览中独立于建模与求图的输入侧支柱。

两种 Dataset 协议

协议	需实现	随机访问	适用场景
map-style	`__len__` + `__getitem__`	支持	数据可索引：本地文件、HDF5 数据格式、内存表
iterable-style	`__iter__`	不支持	流式数据：网络、日志、TAR 分片（WebDataset）

from torch.utils.data import Dataset, IterableDataset
 
class FileDS(Dataset):          # map-style：sampler 生成下标，逐条取
    def __len__(self): ...
    def __getitem__(self, i): ...
 
class StreamDS(IterableDataset):  # iterable-style：只能顺序产出
    def __iter__(self):
        info = get_worker_info()
        for i, x in enumerate(self.src):
            # 多 worker 时各取一片，否则每个 worker 都会拿到全量重复数据
            if info is None or i % info.num_workers == info.worker_id:
                yield x

选型：能给出长度就用 map-style——它才能 shuffle、按需下采样、配合 sampler 切分；iterable-style 放弃 __len__ 与全局乱序，换来对无限流和远端存储的适配。

DataLoader 参数全解

DataLoader(ds, batch_size=32, shuffle=True, num_workers=4,
           pin_memory=True, persistent_workers=True, prefetch_factor=2)

参数	默认	作用
`batch_size`	1	每批样本数
`shuffle`	False	每 epoch 重排下标；仅 map-style，与 `sampler` 互斥
`sampler` / `batch_sampler`	None	自定义取样策略，如 `DistributedSampler`
`drop_last`	False	丢弃不足一个 batch 的尾部；BN/多 rank 对齐时常用 True
`num_workers`	0	加载子进程数，0 表示主进程同步加载
`collate_fn`	`default_collate`	样本列表 → batch 的组装逻辑
`pin_memory`	False	批数据拷入页锁定内存
`prefetch_factor`	2	每个 worker 预取的 batch 数（仅 `num_workers>0` 生效）
`persistent_workers`	False	epoch 间保留 worker 进程，免去重复 fork 与 dataset 重建
`worker_init_fn`	None	worker 启动回调，典型用途是派生随机种子

调优组合：GPU 训练开 pin_memory=True，配合 .to(device, non_blocking=True) 才能异步 H2D 拷贝（见 torch.cuda 与混合精度）；数据集小、epoch 短时开 persistent_workers=True 省下每轮进程启动；num_workers 从 4–8 起调，GPU 利用率仍波动再加大——超过 I/O 吞吐后继续加只会推高内存与进程开销。瓶颈定位先看 nvidia-smi：GPU 利用率锯齿状即数据侧饥饿。

多进程加载原理与坑

num_workers>0 时主进程把 dataset 连同参数 pickle 后复制给各 worker；worker 从共享下标队列领任务、加载并 collate，成品 batch 经共享内存回传。三个高频坑都源于这个架构：

随机性重复：torch 会自动给每个 worker 派生不同种子，但 random 与 NumPy 不会，必须手动处理：

import numpy as np, torch
 
def seed_worker(worker_id):
    s = torch.initial_seed() % 2**32   # DataLoader 已按 worker 错开的 torch 种子
    np.random.seed(s)
    import random; random.seed(s)
 
loader = DataLoader(ds, num_workers=4, worker_init_fn=seed_worker,
                    generator=torch.Generator().manual_seed(42))

序列化失败：macOS/Windows 默认 spawn 启动 worker，dataset、collate_fn、worker_init_fn 走完整 pickle——lambda、局部闭包、已打开的文件句柄与数据库连接都会报错。全部收敛到模块级可 pickle 对象，或改成惰性打开（下文 HDF5 示例）。
fork 与 CUDA 冲突：Linux 默认 fork，但主进程已初始化 CUDA 上下文时 fork 会崩溃；用 CUDA 又要多进程，设置多进程启动方式为 spawn 最稳。另注意 dataset 被整体复制，索引大表内存占用 ×num_workers。

collate_fn 自定义

默认 default_collate 把样本逐字段堆叠成 Tensor 与 Storage 支撑的批张量（dict 按 key 合并、数值转 Tensor、字符串退化为 list）。变长样本必须接管，典型如 padding：

from torch.nn.utils.rnn import pad_sequence
 
def collate(batch):
    xs, ys = zip(*batch)                     # list of (变长 seq, label)
    return pad_sequence(xs, batch_first=True), torch.tensor(ys)
 
loader = DataLoader(ds, batch_size=32, collate_fn=collate)

分布式：DistributedSampler

多卡数据并行要求各 rank 拿到互不重叠的子集，这由 sampler 而非 DataLoader 完成，整体衔接见分布式训练总览：

from torch.distributed import DistributedSampler
 
sampler = DistributedSampler(ds, shuffle=True, drop_last=True)
loader = DataLoader(ds, batch_size=32, sampler=sampler)  # 不再传 shuffle
for epoch in range(epochs):
    sampler.set_epoch(epoch)   # 不调用则每个 epoch 划分方式完全相同
    ...

要点：sampler 会替换 shuffle，两者互斥；sampler 以重复采样补齐使各 rank 样本数对齐，要求 exactly-once 就开 drop_last；set_epoch 混入 epoch 改变排列，漏掉是分布式训练经典 bug。rank 与进程组的来源在 torchrun 启动脚本里就绪。

不用 sampler 的死锁：若各 rank 独立加载全量数据、未用 sampler 切分，数据集长度不整除时各 rank 的 batch 数不等，先跑完的 rank 退出训练循环，其余 rank 卡在梯度 all-reduce 上无限等待——这是 DDP 最典型的挂死方式。DistributedSampler 从源头根治：drop_last=False 以重复采样 pad 至各 rank 等量，drop_last=True 丢弃尾部至等量，两种模式下 rank 间样本数恒相等；DataLoader.drop_last 只作用于单卡内部组 batch，修复不了 rank 间不均。

与存储方案的配合

HDF5 数据格式：单文件随机访问，与 map-style 天作之合；但句柄不能跨进程序列化，须在 __getitem__ 内惰性打开：

class H5DS(Dataset):
    def __init__(self, path):
        self.path, self.f = path, None
        with h5py.File(path, "r") as f:      # 只在主进程读元信息
            self.n = f["x"].shape[0]
    def __len__(self):
        return self.n
    def __getitem__(self, i):
        if self.f is None:                   # 各 worker 首次访问时自行打开
            self.f = h5py.File(self.path, "r")
        return self.f["x"][i], self.f["y"][i]

WebDataset：TAR 分片顺序流式读，对应 iterable-style；分片即并行单位——一个 worker 读一个 shard，天然扩展到对象存储与千级数据集，代价是放弃随机访问与精确 epoch 边界。本地小数据用前者，云上大数据用后者。

知识花园

探索