4. 数据加载参数配置（二）：collate_fn、sampler、worker_init_fn、prefetch_factor

上一节我们把 DataLoader 的基础参数捋了一遍。今天聊几个进阶但极其关键的参数——collate_fn、sampler、worker_init_fn 和 prefetch_factor。这四个家伙，玩不好会掉坑，玩好了能让训练速度和稳定性上一个台阶。

先看一张关系图，帮你快速建立整体认知。

DataLoader 核心参数协作流程 Dataset 原始样本 + 标签 Worker 进程 (num_workers > 0 时启用) worker_init_fn 在此执行 → 自定义每个 worker 的初始化 每个 worker 独立调用 Dataset.__getitem__，并行读取 预取队列 prefetch_factor 控制队列长度 Sampler 控制样本顺序和划分 collate_fn 将 batch 样本合并成 tensor 主干线：Dataset → Worker(并行预处理) → 预取队列 → Sampler索引采样 → collate_fn组batch

4.1 collate_fn：组 batch 的万能接口

collate_fn 是 DataLoader 中最容易被忽视但威力巨大的参数。默认情况下，PyTorch 会把从 Dataset 里取出的样本列表自动堆叠成 batch。但现实中的样本长得五花八门——有图像、有文本、有变长序列、有字典结构，默认行为经常直接炸给你看。

核心原则：collate_fn 接收一个样本列表（list of samples），返回一个 batch。只要你的数据需要组合、填充、遮掩或重新格式，就该写自定义 collate_fn。

什么时候必须用？

  • 变长序列需要 padding（如 NLP 任务）。
  • 样本中包含多个字段（比如 image + mask + metadata）。
  • 需要做数据增强后的重采样或归一化统计。
  • 混合精度训练时需要确保 batch 内数据类型统一。

我举个例子。假设你的 Dataset 返回的是字典：

def __getitem__(self, idx):
    return {
        'image': img_tensor,
        'label': label,
        'filename': self.files[idx]
    }

默认的 collate_fn 遇到字典就会抛异常。你得自己写：

def my_collate(batch):
    images = torch.stack([item['image'] for item in batch], dim=0)
    labels = torch.tensor([item['label'] for item in batch])
    filenames = [item['filename'] for item in batch]
    return {
        'image': images,
        'label': labels,
        'filename': filenames
    }
个人习惯：我通常在 collate_fn 里顺手做「统一 shape 检查」和「nan/inf 检验」。早发现早止损，别等到反向传播才暴露。

掩码 padding 的经典写法（以 Transformer 训练场景为例）：

def collate_transformer(batch):
    input_ids = [item['input_ids'] for item in batch]
    # 计算每个序列长度
    lengths = [len(ids) for ids in input_ids]
    max_len = max(lengths)
    # 填充到 max_len
    padded = torch.zeros(len(batch), max_len, dtype=torch.long)
    attention_mask = torch.zeros(len(batch), max_len, dtype=torch.long)
    for i, ids in enumerate(input_ids):
        padded[i, :len(ids)] = torch.tensor(ids)
        attention_mask[i, :len(ids)] = 1
    labels = torch.tensor([item['label'] for item in batch])
    return padded, attention_mask, labels
我曾经踩过的坑：在 collate_fn 里做随机数据增强（比如随机裁剪）。多个 worker 各自调用 collate_fn，导致同一个数据在不同 epoch 里随机性不一致或产生死锁。后来我把增强移到 Dataset 的 __getitem__ 里，collate_fn 只做纯拼接，问题解决。

4.2 sampler：控制样本出场顺序

sampler 决定了 DataLoader 以什么样的顺序从 Dataset 里抽取样本。默认是顺序抽取，但很多场景需要自定义顺序。

常见 sampler 类型：

Sampler用途典型场景
SequentialSampler按索引顺序逐个取验证集、测试集
RandomSampler随机打乱（默认）训练集（必须 shuffle）
SubsetRandomSampler从子集索引中随机取划分验证集
WeightedRandomSampler按权重采样（可放回）类别不均衡
BatchSampler按 batch 维度包装其他 sampler自定义 batch 内样本组成
DistributedSampler分布式训练，每个进程取不同子集DDP 多卡训练

WeightedRandomSampler 的实战用法：

from torch.utils.data import WeightedRandomSampler

# 假设 labels 是长度为 N 的列表，0～9 共 10 类
class_counts = np.bincount(labels)
weights = 1.0 / class_counts[labels]  # 逆频率加权
sampler = WeightedRandomSampler(weights, num_samples=len(weights), replacement=True)
loader = DataLoader(dataset, batch_size=64, sampler=sampler)

这么做能让模型多看样本少的类别。我以前做医疗图像分类时，正样本占比不到 1%，用这个后 F1 从 0.22 直接跳到 0.67。

注意：shuffle 参数和 sampler 是互斥的——如果传了 sampler，shuffle 必须为 False。DataLoader 内部会弹 warning，我建议代码里显式写明 shuffle=False 降低困扰。

4.3 worker_init_fn：多进程的初始化钩子

当 num_workers > 0 时，每个 worker 子进程都会独立复制 Dataset 并调用 __getitem__。但有些操作需要在每个 worker 启动时只做一次——比如设置随机种子、加载大型缓存文件、建立数据库连接。

这就是 worker_init_fn 的用途。它是一个 callable，签名如下：

def worker_init_fn(worker_id):
    # worker_id 从 0 开始编号
    np.random.seed(worker_id + 42)
    random.seed(worker_id + 42)
    # 如果使用 PyTorch 的 DataLoader，torch 种子会自动设好
灵魂拷问：为什么每个 worker 要单独设随机种子？
因为如果不设，所有 worker 可能会共享主进程的随机状态，导致多个 worker 生成一模一样的增强结果，明明参数 shuffle 了却看到大量重复样本。嗯，我有个同事.debug 了一整天才找到原因。

更实用的场景：共享缓存

def worker_init_fn(worker_id):
    global shared_cache
    if shared_cache is None:
        # 只在第一个 worker 里加载一次，然后存在全局变量
        shared_cache = load_big_file('precomputed_embeddings.npy')
    # 每个 worker 都可以访问 shared_cache

注意：跨进程共享全局变量要小心，Python 的 multiprocessing 默认是用 fork 或 spawn，worker 之间内存不共享，要用共享内存或 mmap。

我曾经遇到过的怪问题：在 worker_init_fn 里连接数据库，结果 worker 数量一多，数据库连接池打满导致训练卡死。后来改成每个 worker 启动时只连接一次，并用 worker_init_fn 控制连接复用，数据库端也加大了连接数上限。

4.4 prefetch_factor：预取队列的油门

这个参数控制每个 worker 在后台预取多少个 batch 存入队列。prefetch_factor=2 表示每个 worker 提前准备 2 个 batch 等主进程来取（默认值也是 2）。

它本质上是个吞吐量 vs 内存的权衡：

  • 预取越多 → GPU 等数据的概率越低 → 吞吐提升。
  • 预取越多 → 内存中暂存的 batch 越多 → 显存/内存压力增大。

我来拆一下数学关系。总的预取 batch 数为：

total_prefetch_batches = num_workers * prefetch_factor

如果你的 batch 很大（比如每样本 256 张图），num_workers=8、prefetch_factor=2 就相当于内存里同时放着 16 个 batch 的数据。稍不注意就把内存撑爆。

最佳实践建议：

  1. 内存充足：prefetch_factor=4 甚至 8，最大化吞吐。
  2. 内存紧张：设回 2 或 1，也可以尝试 prefetch_factor=0（禁用预取，但会增加等待）。
  3. 显存不均衡：配合 pin_memory=True 一起用，预取的内存页锁定，传输到 GPU 更快。
一个调优经验：我曾在一个视频理解任务中，prefetch_factor=2 时 GPU 利用率只有 70%，调成 4 后升到 95%，内存从 32GB 涨到 48GB——但稳住了。如果你的 GPU 利用率一直上不去，先查 I/O 是不是瓶颈，然后调这个值。

4.5 这四个参数怎么协同工作？

拿一个典型的多卡分布式训练场景来说：

from torch.utils.data import DataLoader, DistributedSampler
from torch.utils.data.distributed import DistributedSampler

sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)
loader = DataLoader(
    dataset,
    batch_size=32,
    sampler=sampler,
    num_workers=4,
    prefetch_factor=2,
    collate_fn=my_custom_collate,
    worker_init_fn=worker_seed_fn,
    pin_memory=True
)

流程是这样的：

  1. DistributedSampler 告诉 DataLoader 只取属于当前 rank 的样本索引。
  2. 4 个 Worker 各自启动，worker_init_fn 设好种子和连接。
  3. 每个 Worker 从 Dataset 里取 2 个 batch 预取进队列（prefetch_factor=2）。
  4. 主进程消费时，Worker 把样本列表传给 collate_fn 组 batch。
  5. 每个 batch 被 pin_memory 固定后送到 GPU。
调试小技巧：如果在 collate_fn 发生异常，错误栈往往指向 worker 内部，不好定位。我习惯先在 num_workers=0 下跑一个小规模数据测试 collate_fn，确认无误后再开多进程。

这四个参数配合得好，你可以做到：

  • 任意复杂的 batch 结构（collate_fn）。
  • 灵活控制样本顺序和重采样（sampler）。
  • 多进程稳定运行（worker_init_fn）。
  • I/O 几乎不阻塞训练（prefetch_factor）。

说到底，数据加载是训练 Pipeline 里最容易忽视的瓶颈，也是优化空间最大的环节。这几把钥匙给你了，用好了能省下你不少调参时间。

← 上一章 📖 返回目录 下一章 →