数据加载优化：别让GPU饿着等数据

做深度学习训练，我最怕看到一种情况：GPU利用率像心电图一样忽高忽低。明明显存够用，计算能力也强，但就是跑不快。问题出在哪？十有八九是数据加载拖了后腿。

说白了，GPU就像一台高速运转的机器，而数据加载就是给这台机器喂原料的传送带。传送带慢了，机器再快也没用。今天我们就聊聊怎么把这条传送带调顺溜。

DataLoader工作原理：数据是怎么送到GPU的？

先搞清楚DataLoader到底在干什么。它的核心任务就一个：把硬盘上的数据，变成GPU能吃的张量。

流程大概是这样的：

从硬盘读取原始数据（图片、文本等）
做预处理（缩放、归一化、数据增强）
打包成batch
送到GPU显存里

嗯，听起来简单。但这里有个关键点——这些操作是串行的还是并行的？

默认情况下，DataLoader是在主进程里一个一个处理的。你想想看，GPU在疯狂计算，主进程却在等硬盘读数据，这不就浪费了吗？

核心瓶颈：数据加载和模型训练天然是异步的。如果加载跟不上训练，GPU就得空转等着。

num_workers调优：多几个搬运工

解决这个问题最直接的办法——多开几个子进程来干活。这就是num_workers参数的作用。

我个人的习惯是，先设成CPU核心数的一半试试。比如你机器有8核，就先设4。为什么不是越多越好？

我在项目中遇到过这种情况：把num_workers设成16，结果训练反而变慢了。排查了半天，发现是进程间通信开销太大，数据拷贝占用了大量内存带宽。

调优建议：

起步值：CPU核心数的一半
观察GPU利用率，如果低于80%，逐步增加
上限：不要超过CPU核心数
注意内存占用，每个worker都会复制一份数据

这里有个小技巧：你可以用nvidia-smi和htop同时观察GPU和CPU的状态。如果GPU利用率低但CPU没跑满，说明瓶颈不在数据加载；如果CPU跑满了GPU还在等，那就继续加worker。

prefetch_factor设置：提前准备下一批

光有多个worker还不够，还得让它们提前干活。这就是prefetch_factor的作用。

默认情况下，DataLoader是"懒加载"的——当前batch用完了，才去准备下一个。这就像你吃完饭才去洗碗，中间总有个空档期。

prefetch_factor就是让worker提前准备几个batch的数据。默认值是2，意味着每个worker会提前准备2个batch。

# 设置prefetch_factor
DataLoader(
    dataset,
    batch_size=32,
    num_workers=4,
    prefetch_factor=4,  # 每个worker提前准备4个batch
    persistent_workers=True  # worker进程不销毁，复用
)

我记得有一次调优，把prefetch_factor从2调到8，训练速度提升了15%。但要注意，这个值不是越大越好——它会占用更多内存来缓存数据。

注意：prefetch_factor和num_workers是联动的。总预取数量 = num_workers × prefetch_factor。设太大可能导致内存溢出。

内存映射(mmap)文件读取：让操作系统帮你干活

前面说的都是多进程层面的优化。但还有一个更底层的优化——文件读取方式。

传统读取文件，是先把数据从硬盘拷贝到内核缓冲区，再拷贝到用户空间。这中间有两次拷贝，很费时间。

内存映射（mmap）就不一样了。它直接把文件映射到进程的虚拟地址空间，操作系统按需加载数据。说白了，就是让操作系统帮你管理文件缓存，省去了手动拷贝的步骤。

# 使用mmap加载数据
import numpy as np
import mmap

# 传统方式
data = np.fromfile('data.bin', dtype=np.float32)

# mmap方式
with open('data.bin', 'rb') as f:
    mmapped = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ)
    data = np.frombuffer(mmapped, dtype=np.float32).reshape(-1, 3, 224, 224)

我曾经在一个图像分类项目里，把数据加载时间从每epoch 45秒降到了12秒。怎么做到的？就是把所有图片打包成一个二进制文件，然后用mmap读取。

mmap的优势：

减少数据拷贝次数
利用操作系统的页面缓存
多个进程可以共享同一份内存映射
按需加载，不用一次性读入全部数据

但要注意，mmap不是银弹。如果你的数据是随机访问的，而且访问模式很分散，mmap反而可能因为频繁缺页中断而变慢。

实战配置模板

说了这么多，给一个我常用的配置模板：

from torch.utils.data import DataLoader

def create_dataloader(dataset, batch_size, num_workers=None):
    # 自动检测CPU核心数
    cpu_count = os.cpu_count()
    if num_workers is None:
        num_workers = min(4, cpu_count // 2)
    
    return DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=True,
        num_workers=num_workers,
        prefetch_factor=4,
        persistent_workers=True,
        pin_memory=True,  # 锁页内存，加速CPU到GPU传输
        drop_last=True    # 丢弃最后一个不完整的batch
    )

这个配置在我经手的项目里，基本都能把GPU利用率稳定在90%以上。当然，具体数值还得根据你的硬件和数据集来调。

知识体系总览

下面这张图总结了数据加载优化的核心逻辑：

一句话总结：数据加载优化的目标就是让GPU永远有活干。用多进程并行加载，用预取机制提前准备，用mmap减少IO开销。这三板斧下去，大部分场景都能搞定。

想要更多资料，联系可可

📚 课程目录

数据加载优化：别让GPU饿着等数据

DataLoader工作原理：数据是怎么送到GPU的？

num_workers调优：多几个搬运工

prefetch_factor设置：提前准备下一批

内存映射(mmap)文件读取：让操作系统帮你干活

实战配置模板

知识体系总览