4. 数据加载优化：DataLoader 工作原理、num_workers 设置、pin_memory 的作用、prefetch_factor 调整

聊到显存溢出，很多人第一反应是模型太大、batch size 太高。但说实话，我排查过不少案例，最后发现罪魁祸首其实是数据加载那一块。数据没喂好，GPU 就得干等着，或者一下子把显存撑爆。今天咱们就把 DataLoader 这层皮扒开，看看里面到底怎么回事。

4.1 DataLoader 的工作原理：它到底在忙什么？

DataLoader 说白了就是一个「数据搬运工」。它的任务很简单：从硬盘上把数据读进来，做预处理（比如裁剪、归一化），然后打包成 batch 送给 GPU。

但这里有个关键点——它是一次性把所有数据都塞进内存吗？不是的。它内部维护了一个迭代器，每次只加载当前 batch 需要的数据。你想想看，如果你的数据集有 100 万张图片，一次性全加载进来，内存早炸了。

我个人的习惯是，把 DataLoader 想象成一个「流水线」：

  • 第一步：从硬盘读取原始数据（比如图片路径、标签）
  • 第二步：执行用户定义的 transform（比如随机翻转、归一化）
  • 第三步：将处理好的数据组装成一个 batch 的 tensor
  • 第四步：把这个 batch 送到 GPU 上

嗯，这里要注意：如果 transform 里做了很重的计算（比如大尺寸图像的随机裁剪），那这一步就会成为瓶颈。我曾经遇到过一个项目，transform 里用了复杂的图像增强库，结果 CPU 忙得冒烟，GPU 却闲得发慌。显存没溢出，但训练速度慢得像蜗牛。

4.2 num_workers 设置：多线程到底开几个？

num_workers 这个参数，很多人直接设成 4 或者 8，觉得越多越好。其实不是这样的。

它的作用是什么呢？就是让多个子进程同时去加载数据。比如你设 num_workers=4，那 DataLoader 会启动 4 个进程，每个进程独立地从硬盘读数据、做预处理。这样主进程（也就是训练循环）就不用等数据加载了。

但这里有个坑：不是 workers 越多越好。我见过有人把 num_workers 设成 32，结果训练反而变慢了。为什么会这样？

  • CPU 瓶颈：每个 worker 都要占用 CPU 资源。如果你的 CPU 核心数只有 8 个，你开 32 个 worker，那大部分时间都花在进程切换上了。
  • 内存压力：每个 worker 都会复制一份数据到自己的内存空间。worker 越多，内存占用越大。本来显存就紧张，内存再被吃光，系统开始用 swap，那训练直接卡死。
  • IO 竞争：多个 worker 同时读硬盘，如果硬盘速度跟不上（比如机械硬盘），反而会互相干扰。
我的经验值：一般建议 num_workers 设为 CPU 核心数的一半或等于核心数。比如 8 核 CPU，设 4 或 8。如果数据预处理很轻（比如只是读图片、转 tensor），设 2 就够。如果预处理很重（比如做大量数据增强），可以适当增加。

另外，如果你在 Windows 上训练，num_workers 设成 0 反而更稳定。Windows 的多进程实现有点特殊，容易出 bug。我踩过这个坑，后来干脆在 Windows 上只用单进程。

4.3 pin_memory 的作用：给 GPU 铺一条高速公路

pin_memory 这个参数，很多人直接设成 True 就完事了。但你真的理解它在干什么吗？

简单说，pin_memory 就是把数据固定在内存的某个特殊区域。这个区域叫「页锁定内存」（pinned memory）。普通的内存，操作系统可以随时把数据挪来挪去（比如为了给其他程序腾空间）。但页锁定内存不会，它一直待在那里不动。

为什么这对 GPU 有好处？因为 GPU 从 CPU 拿数据时，需要经过 PCIe 总线。如果数据在普通内存里，GPU 得先通过 CPU 把数据复制到页锁定内存，然后再传输。多了一步操作，速度就慢了。

如果数据一开始就在页锁定内存里，GPU 可以直接通过 DMA（直接内存访问）把数据拉过去，省掉了中间步骤。说白了，就是给 GPU 铺了一条高速公路。

关键点：pin_memory=True 会占用更多的系统内存（因为页锁定内存不能被换出到硬盘）。如果你的系统内存本身就不够（比如只有 8GB），设成 True 反而可能导致内存不足。我建议：内存 >= 16GB 时放心开，内存 < 8GB 时谨慎使用。

我记得有一次帮同事排查显存溢出，他开了 pin_memory=True，但系统内存只有 8GB，结果数据加载时内存爆了，连带显存也报 OOM。关掉 pin_memory 后，问题就解决了。所以这个参数不是无脑开的。

4.4 prefetch_factor 调整：提前准备，别让 GPU 饿着

prefetch_factor 是 PyTorch 1.7 之后引入的参数。它的作用是控制 DataLoader 提前加载多少个 batch 的数据。

默认情况下，prefetch_factor=2，意思是 DataLoader 会提前加载 2 个 batch 的数据到内存里。这样当 GPU 正在处理当前 batch 时，下一个 batch 已经在内存里等着了，GPU 不用等。

这个参数在什么场景下有用？

  • 数据加载慢：如果你的 transform 很重，或者硬盘速度慢，GPU 经常处于等待状态。这时候增大 prefetch_factor（比如设成 4 或 8），可以让 DataLoader 提前多准备几个 batch，减少 GPU 空闲时间。
  • 显存紧张：但要注意，prefetch_factor 越大，内存里缓存的 batch 就越多。如果每个 batch 很大（比如高分辨率图像），这些缓存会占用大量内存。内存不够时，系统会开始用 swap，反而更慢。
避坑指南：我曾经把 prefetch_factor 设成 16，想着多准备点总没错。结果内存直接飙到 30GB，系统卡死。后来发现，每个 batch 是 16 张 4K 图像，16 个 batch 就是 256 张图，内存根本扛不住。所以这个参数要根据 batch size 和图像大小来调，别贪心。

一般来说，我建议这样调：

  • 如果 GPU 利用率经常低于 80%，且 CPU 没有跑满，可以增大 prefetch_factor。
  • 如果内存占用已经很高（比如超过 80%），就减小 prefetch_factor 或者保持默认。
  • 如果数据加载很快（比如 SSD + 轻量 transform），默认的 2 就够用了。

4.5 知识体系总览

下面这张图总结了 DataLoader 优化的核心逻辑。你可以看到，数据从硬盘到 GPU 的每一步都可能成为瓶颈，而 num_workers、pin_memory、prefetch_factor 就是用来打通这些瓶颈的工具。

DataLoader 优化核心逻辑 硬盘 读取 num_workers 多进程并行加载 预处理 transform 裁剪/归一化等 缓存 prefetch_factor 预取 batch 数 传输 pin_memory 页锁定内存加速传输 PCIe GPU 显存 优化目标：减少 GPU 等待时间，避免内存/显存溢出 每个环节都可能成为瓶颈，需要根据硬件配置动态调整

4.6 实战建议：一套组合拳

说了这么多理论，最后给一套我常用的配置模板。你可以根据自己的硬件情况微调：

from torch.utils.data import DataLoader

# 假设你的 CPU 是 8 核，内存 32GB，GPU 显存 8GB
dataloader = DataLoader(
    dataset,
    batch_size=32,
    shuffle=True,
    num_workers=4,          # CPU 核心数的一半
    pin_memory=True,        # 内存够用就开
    prefetch_factor=4,      # 预取 4 个 batch
    persistent_workers=True # 保持 worker 进程存活，减少创建开销
)

这里我加了 persistent_workers=True，这个参数的作用是让 worker 进程在 epoch 之间不销毁，避免每次重新创建进程的开销。如果你的数据集不大（比如几万张图），这个参数效果不明显。但如果是百万级数据集，它能省下不少时间。

一个小技巧：如果你不确定 num_workers 该设多少，可以写个简单的测试脚本，循环跑几个 epoch，观察 GPU 利用率。如果 GPU 利用率经常低于 70%，就增加 num_workers 或 prefetch_factor。如果内存占用接近 90%，就减小它们。

嗯，数据加载优化这块，说白了就是「让 GPU 别闲着，也别撑爆」。每个参数都有它的适用场景，没有万能配置。我建议你每次调参时只改一个变量，观察效果，这样才知道哪个参数起了作用。

最后提醒一句：如果你用的是分布式训练（多卡），每个 GPU 都会有自己的 DataLoader。这时候 num_workers 要乘以 GPU 数量来算总进程数，别让 CPU 过载了。

可可精品资料库

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