📚 课程目录第4章数据加载与预处理加速

4. 数据加载与预处理加速：让GPU不再干等

你有没有遇到过这种情况？
GPU利用率在50%上下波动，显存占用忽高忽低。
训练一个epoch，真正算的时间还没等数据的时间长。
这时候别急着换显卡，问题很可能出在数据加载这个环节上。

我刚开始做大模型训练时，就踩过这个坑。
V100跑着跑着就“喘气”，一看监控，GPU计算时间只占30%。
剩下70%都在等CPU把数据送过来。说白了，就是饭没端上桌。

本章咱们重点聊三个加速数据加载的手段：
DataLoader的num_workers、pin_memory，以及NVIDIA的杀手锏——DALI库。

4.1 先说数据加载的瓶颈在哪

单线程加载一个batch的流程：

读文件（磁盘I/O）
解码（JPEG/PNG）
数据增强（随机裁剪、翻转、归一化）
转Tensor
拷贝到GPU

这一套下来，如果是高分辨率图像，单步可能耗费几百毫秒。
GPU算一个batch可能只需要10毫秒。你说它不等还能干嘛？

核心原则： 数据加载管线要把“喂数据”的速度，匹配上“消化数据”的速度。

4.2 num_workers：偷懒就靠多线程

num_workers是PyTorch DataLoader最直观的参数。
它决定了有几个子进程在后台帮你准备数据。

num_workers	特征	我个人习惯
0	主进程一个线程干所有事，慢	小数据集调试时用
1	一个子进程加载，稍微好点	几乎不用
2~4	常用区间，性价比高	CPU核数的一半左右
8~16	高并发，适合高频预处理	大数据集训练时推荐

from torch.utils.data import DataLoader

dataloader = DataLoader(
    dataset,
    batch_size=32,
    shuffle=True,
    num_workers=4,      # 根据CPU核心数调整
    pin_memory=True
)

这里有个坑我不得不说。
num_workers不是越大越好。
我曾经把一台48核机器设成48，结果训练直接跑崩了。
为什么？子进程多了，争抢文件句柄和内存带宽，反而互相拖慢。

经验法则： num_workers设为CPU物理核心数的一半，或（核心数-2）。再多，边际收益就递减了。

4.3 pin_memory：给GPU铺条高速公路

pin_memory=True这行代码，很多人不知道具体做了啥。
其实就是把CPU上的数据锁在物理内存里，不让操作系统把它换到磁盘交换区。

为什么要锁住？
GPU要通过PCIe总线拷贝数据。
如果数据在物理内存里到处乱漂（虚拟内存机制），DMA控制器根本找不到它，只能走CPU中转，慢得离谱。

我上一个项目，不加pin_memory时，一次拷贝耗时2ms。加上后降到0.1ms。
20倍的差距，成本只是多占用一点固定内存，超值。

# 推荐配置
train_loader = DataLoader(
    dataset,
    batch_size=64,
    shuffle=True,
    num_workers=4,
    pin_memory=True      # 打开，除非内存不够
)

注意： pin_memory会在物理内存中锁定一块区域。如果内存吃紧，系统可能会变卡。一般建议总锁定量不超过物理内存的30%。

4.4 NVIDIA DALI：工业级数据预处理引擎

前面两个方法能解决一部分问题。
但如果数据预处理特别重（比如视频帧、医疗影像、3D点云），PyTorch自带的DataLoader还是不够快。

这时候就该DALI上场了。

DALI是NVIDIA开源的库，全称Data Loading Library。
它把数据预处理直接放到GPU上跑，绕过CPU瓶颈。

说白了，DALI干了两件事：

异步流水线：数据加载、解码、增强全部流水线化，不阻塞主线程
GPU算子：裁剪、颜色变换、归一化直接用GPU计算，速度是CPU的5~10倍

4.4.1 一个真实案例

我记得有个做医学影像分割的项目，输入是512x512的CT切片，每张图需要做随机弹性变形、高斯噪声、归一化。
用PyTorch原生的DataLoader，一个batch（32张）预处理耗时约180ms。
换成DALI后，降到18ms。训练速度直接翻了将近3倍。

4.4.2 快速上手DALI

import nvidia.dali as dali
import nvidia.dali.fn as fn
import nvidia.dali.types as types

@pipeline_def(batch_size=64, num_threads=4, device_id=0)
def train_pipeline():
    # 从文件列表读取
    jpegs, labels = fn.readers.file(file_root="./images", random_shuffle=True)
    # 硬件解码（JPEG）
    images = fn.decoders.image(jpegs, device="mixed")
    # 在GPU上做数据增强
    images = fn.resize(images, resize_x=224, resize_y=224)
    images = fn.flip(images, horizontal=fn.random.coin_flip())
    images = fn.crop_mirror_normalize(
        images,
        dtype=types.FLOAT,
        mean=[0.485 * 255, 0.456 * 255, 0.406 * 255],
        std=[0.229 * 255, 0.224 * 255, 0.225 * 255]
    )
    return images, labels

pipeline = train_pipeline()
pipeline.build()

你注意看device="mixed"这一行。
意思是：解码时用CPU（混合模式），后面的增强全在GPU上。
这样CPU只管读文件，GPU专管计算，各司其职。

4.5 三种方案的对比

方案	瓶颈	加速效果	适用场景
num_workers调优	CPU解码 + 预处理	3~5倍	日常训练，简单预处理
+ pin_memory	CPU到GPU拷贝	2~20倍拷贝加速	所有GPU训练场景
+ DALI	CPU预处理 + 拷贝	5~10倍整体加速	高分辨率、重预处理

我的建议： 

小项目：num_workers + pin_memory 就够。

大项目（每秒处理上千张图）：直接上DALI。

4.6 一张图看懂数据加载加速架构

下面这个图，很直观地展示了数据从磁盘到GPU的路径。
以及每个环节的加速手段放在哪里。

4.7 一些更具体的避坑指南

内存泄漏：num_workers设太高，每个worker都会复制一份数据集对象。我之前用torch.multiprocessing踩过这个坑，显存撑爆了才发现。解决办法：用fork模式或设置worker_init_fn来控制资源。
随机种子不一致：多worker时每个进程的随机种子要是独立的，否则增强后的数据都一样。PyTorch 1.9以后默认处理了，老版本需要手动设worker_init_fn传种子。
文件描述符耗尽：假如你的数据是几百万个小文件，num_workers一多，每个worker都要打开文件句柄，可能超过系统限制。我那时用 ulimit -n 65536 解决的。

一个实用checklist：
☑ 先用 num_workers=0 跑一轮，观察GPU利用率。
☑ 调高 num_workers（核心数一半），加 pin_memory=True。
☑ 如果预处理重，再评估是否上DALI。
☑ 监控系统内存和GPU利用率，别让数据加载成为墙。

嗯，数据加载这一环，很多人觉得简单就忽略了。
但实际上，它往往是整个训练管线里最容易被卡住的点。
把这块优化好，训练速度提升个2~3倍，比我换一张显卡划算得多。