PyTorch DataLoader 怎样控制多进程预取内存

我在把 PyTorch 训练任务从单进程切到多进程时，最容易误判的一点是：num_workers 变大后，内存并不是只增加几个 Python 进程的开销。DataLoader 还会为每个 worker 准备预取批次，批次越大、worker 越多，主机侧同时挂着的数据就越多。

官方地址：https://pytorch.org/docs/stable/data.html

控制这部分内存，最直接的抓手是把 num_workers 和 prefetch_factor 一起调小，再根据吞吐量决定是否保留 pin_memory=True。一个实用的估算是：

预取中的批次数约等于 num_workers × prefetch_factor，预取数据的主机内存约等于这个批次数乘以单个 batch 的字节数；实际占用还要加上 worker 副本、解码缓存、collate 临时对象和 pinned memory 的额外空间。

官方文档说明，prefetch_factor 表示每个 worker 提前加载的 batch 数；它不是所有 worker 合计的值。本文只讨论这个预取内存问题，不把它混同为模型参数或 CUDA 显存管理。

先看清 DataLoader 的内存组成

当 num_workers=0 时，数据加载发生在主进程中，不存在 worker 预取队列。设置为正数后，DataLoader 会创建对应数量的 worker，每个 worker 都可能持有待交付的 batch。对我来说，排查内存上涨时先把下面几项分开，比直接把 worker 数改成 1 更容易找到原因：

• 单批大小：batch_size、样本字段数量和 dtype 决定一批数据有多大。
• 预取批数：num_workers × prefetch_factor 决定多进程阶段可能同时准备多少批。
• 处理临时对象：图片解码、增强和 collate_fn 可能在形成 batch 前制造额外峰值。
• 固定页内存：pin_memory=True 方便向 GPU 传输，但它不是“免费缓存”，仍然需要主机内存。

用公式估算，而不是盲目增加 worker

假设一个 batch 由图像张量、标签和少量元数据组成，先估出单批字节数 batch_bytes。如果使用 4 个 worker、每个 worker 预取 2 个 batch，预取队列就可能容纳约 8 个 batch。单批 64 MB 时，仅预取数据就可能接近 512 MB；这还没算 worker 内的解码中间结果。

下面的代码是估算器，不会读取真实数据，也不代表 PyTorch 会精确按照这个数分配内存。它适合在调整参数前快速判断数量级：

from dataclasses import dataclass

@dataclass
class LoaderPlan:
    # 这些字段只描述 DataLoader 的关键内存旋钮
    batch_bytes: int
    num_workers: int
    prefetch_factor: int

    def prefetched_bytes(self) -> int:
        # 多进程时，每个 worker 都按 prefetch_factor 准备 batch
        return self.batch_bytes * self.num_workers * self.prefetch_factor


plan = LoaderPlan(
    batch_bytes=64 * 1024 * 1024,  # 示例：每个 batch 约 64 MiB
    num_workers=2,                 # 先用较少 worker 观察吞吐和内存
    prefetch_factor=1,             # 每个 worker 只提前准备 1 个 batch
)

print(f"约 {plan.prefetched_bytes() / 1024 ** 2:.0f} MiB 处于预取阶段")
# 这是容量估算，不包含解码、collate 和 pinned memory 的临时开销

需要注意两个边界：第一，prefetch_factor 只在 num_workers > 0 的多进程场景有意义；第二，公式估算的是预取数据，不是整个进程的 RSS。图片增强、字符串列表、缓存对象和数据集副本都可能让实际峰值更高。

三种参数组合，按内存约束做选择

我通常不会先追求最高 worker 数，而是先按机器余量选方案，再看 GPU 是否经常等数据。下面的组合重点是决策逻辑：

平衡方案可以这样写。若 pin_memory=True，它只是改善从主机到 CUDA 设备的搬运条件，并不会减少预取 batch 的数量；内存已经紧张时，应先降低预取批数，再决定是否保留固定页内存。

import torch

loader = torch.utils.data.DataLoader(
    dataset,
    batch_size=32,                 # 先控制单批大小，再增加并发
    shuffle=True,
    num_workers=2,                 # 两个 worker，避免预取规模过快放大
    prefetch_factor=1,             # 每个 worker 只保留一个待交付 batch
    pin_memory=True,               # GPU 训练可尝试；主机内存紧张时先关掉对比
    persistent_workers=False,      # 先不跨 epoch 常驻 worker
)

for batch in loader:
    # non_blocking 需要配合合适的 pinned memory 使用
    inputs, labels = (item.to(device, non_blocking=True) for item in batch)
    loss = train_step(inputs, labels)

这里的 persistent_workers=False 是内存优先的起点。它会让 worker 在迭代结束后退出；如果每个 epoch 创建和销毁 worker 的成本明显，且主机有足够余量，再考虑改成 True。持久 worker 的收益是减少重复启动，代价是数据集实例和相关状态可能跨 epoch 留在内存中。

别漏掉 IterableDataset 的副本放大

如果数据集是 IterableDataset，多进程时每个 worker 会拿到一个数据集副本。没有用 get_worker_info() 或 worker_init_fn 切分范围时，不仅可能重复读取数据，也可能让每个副本都保留自己的文件句柄、缓冲区或解码状态。

因此，降低 prefetch_factor 只能压低待交付 batch 的数量，不能自动消除数据集副本的成本。我的排查顺序通常是：

1. 先把 num_workers 改为 0，确认问题是否来自多进程阶段。
2. 恢复少量 worker，把 prefetch_factor 设为 1，观察内存数量级是否稳定。
3. 如果仍然上涨，检查 IterableDataset 是否按 worker 分片，以及 collate_fn 是否保留了不必要的对象。
4. 只有在内存余量足够时，再尝试 persistent_workers=True 或增加 worker。

最后的决策表

可以把结论压缩成一句话：先控制单批大小，再控制 num_workers × prefetch_factor，最后才用 pin_memory 和持久 worker 换取吞吐。下面是我会留下的速查规则：

• 主机内存紧张：num_workers=0 或少量 worker，优先降低 prefetch_factor。
• GPU 等数据但主机有余量：从 num_workers=2, prefetch_factor=1 开始逐步增加。
• 跨 epoch 启动成本高：确认常驻内存能接受后再用 persistent_workers=True。
• 使用 IterableDataset：先做 worker 分片，否则重复数据和副本缓存会同时放大。

相关问题

prefetch_factor 越大，训练一定越快吗？

不一定。它只增加每个 worker 提前准备的 batch 数；如果瓶颈在磁盘、解码、网络或 GPU 计算，增加预取可能只增加内存，不会带来同等吞吐收益。

pin_memory=True 会减少 CPU 内存吗？

不会。它服务于主机到 GPU 的传输路径，固定页内存仍然是主机内存的一部分。内存优先时，应把它作为单独变量对比，而不是把它当成预取开关。