PyTorch DataLoader 如何限制预取批次避免共享内存耗尽

PyTorch 的 DataLoader 出现共享内存耗尽时，先不要只把 num_workers 调小或调大。真正需要控制的是“同时存在多少个 batch”：worker 数量、每个 worker 的预取批次、单个 batch 的体积会共同决定内存峰值。通常先降低 prefetch_factor 和 batch_size，再按吞吐量逐步增加 worker。

官方文档：https://docs.pytorch.org/docs/stable/data.html

先算清 DataLoader 会预取多少批次

当 num_workers 大于 0 时，数据加载由多个 worker 进程完成。对常见的 map-style 数据集，可以先用一个保守估算理解压力：同时排队的 batch 数 ≈ num_workers × prefetch_factor，而这些 batch 的内存大致还要乘以单个 batch 的大小。这个公式不是精确的内核内存账本，但足以发现配置是否明显超过共享内存预算。

例如，一个 batch 解码后约占 32 MB，设置 8 个 worker、每个 worker 预取 2 个 batch，队列就可能需要约 512 MB 的数据空间；如果样本还包含字符串、临时张量或自定义 collate_fn 的中间对象，实际峰值还会继续上升。先把 prefetch_factor 从默认值降到 1，再观察吞吐量，通常比盲目增加 worker 更容易定位问题。

把共享内存预算写进配置

生产配置应先确定“允许多少 batch 同时在 CPU 侧等待”，再反推参数，而不是把所有机器都套用同一个 worker 数。下面的示例把 worker 和预取值显式写出，注释说明每个参数的边界；代码用于展示配置关系，不把未执行的输出当作运行证据。

from torch.utils.data import DataLoader

# 这个示例把并发预取控制在可估算范围内。
loader = DataLoader(
    dataset,
    batch_size=32,          # 先用较小 batch 控制单个样本包的体积。
    num_workers=4,          # worker 越多，潜在的并行预取槽位越多。
    prefetch_factor=1,      # 每个 worker 只保留一个待处理 batch。
    persistent_workers=True,  # 复用 worker，减少跨 epoch 的重启开销。
    pin_memory=True,        # 为 CPU 到 GPU 的拷贝准备页锁定内存；仍需关注主机内存。
)

for batch in loader:
    # 非阻塞拷贝只有在设备与数据条件满足时才有意义。
    batch = move_batch_to_device(batch, non_blocking=True)
    train_one_batch(batch)

如果单个 batch 约为 B MB，可以把 B × num_workers × prefetch_factor 当作第一版预算。为了解码缓存、框架开销和其他进程留出余量，实际配置不应把这个乘积顶到容器的共享内存上限。共享内存报错消失后，再一次只增加一个变量，记录吞吐量与内存变化。

把 worker 生命周期和内存策略分开配置

persistent_workers=True 的含义是让 worker 在一个迭代器结束后继续存活，适合跨多个 epoch 重复读取且初始化成本较高的数据集。它减少的是进程重新创建和数据集重新初始化的开销，并不会把预取队列变成无限小，也不会自动回收已经排队的 batch。

pin_memory=True 则是另一条边界：它让加载到 CPU 的张量进入页锁定内存，便于向 CUDA 设备发起更高效的拷贝，但页锁定内存仍然属于主机资源。若发现主机内存上涨，应同时检查 prefetch_factor、batch_size、样本解码缓存和 collate_fn，不要只看 GPU 显存。

按现象定位是哪一层超预算

发布前可以固定一份小清单：确认单个 batch 的大致大小；计算 worker 与预取槽位的乘积；给共享内存、页锁定内存和其他进程留余量；确认 persistent_workers 是生命周期选择而非内存限制；每次只改一个参数。这样既能控制共享内存峰值，也能避免为了追求吞吐量把问题转移到主机内存。

常见问题

prefetch_factor=1 还会占用共享内存吗？

会。它只是把每个 worker 的预取槽位降到一个，实际占用仍取决于 worker 数、batch 体积、数据转换和进程间传输开销。

num_workers=0 是否最省内存？

它会取消多进程 worker 队列，通常更容易排查共享内存问题，但数据加载在主进程执行，可能降低训练吞吐量。应把它作为定位手段或小数据集配置，而不是所有场景的默认答案。

persistent_workers=True 能解决 worker 内存泄漏吗？

不能直接解决。它会延长 worker 生命周期；如果数据集、缓存或自定义处理逻辑持续持有对象，问题可能更明显。需要先缩小缓存范围，再决定是否复用 worker。