3、数据加载优化：让GPU不再“饿肚子”

做深度学习训练，最让人头疼的事之一是什么？

我告诉你，不是模型不收敛，而是GPU利用率上不去。

你花大价钱买的A100，结果一看nvidia-smi，GPU利用率才30%。钱白花了，时间也浪费了。这背后最常见的原因，就是数据加载太慢，GPU一直在等数据。

说白了，GPU就像个超级大胃王，你喂食的速度跟不上它吃的速度，它就只能干等着。

今天我们就聊聊怎么把数据喂得快一点。核心就三个参数：num_workers、pin_memory、prefetch_factor。

3.1 num_workers：多几个搬运工

先看一个最简单的例子：

from torch.utils.data import DataLoader

# 默认情况：只有一个进程在加载数据
dataloader = DataLoader(dataset, batch_size=32)

# 优化后：用4个进程并行加载
dataloader = DataLoader(dataset, batch_size=32, num_workers=4)

num_workers这个参数，我习惯叫它“搬运工数量”。默认是0，意思是主进程自己搬数据。你想想看，主进程又要搬数据又要做前向传播，忙得过来吗？

设置成大于0，就会启动子进程专门负责数据加载。GPU在算的时候，子进程已经在准备下一批数据了。

我在项目中遇到过一个问题：把num_workers设成32，结果训练反而变慢了。后来发现是磁盘IO成了瓶颈，太多进程同时读硬盘，互相抢资源。嗯，这里要注意，不是越大越好。

3.2 pin_memory：打通CPU到GPU的“高速公路”

数据从CPU内存搬到GPU显存，这个过程叫数据传输。默认情况下，数据先放在CPU的“普通内存”里，再拷贝到GPU。这就像走乡间小路，慢。

pin_memory=True的作用，是把数据固定在CPU的“锁页内存”里。锁页内存不会被操作系统换到硬盘上，GPU可以直接通过DMA（直接内存访问）读取，速度飞快。

# 普通模式
dataloader = DataLoader(dataset, batch_size=32, num_workers=4)

# 加速模式
dataloader = DataLoader(dataset, batch_size=32, num_workers=4, pin_memory=True)

我曾经在一个OCR识别项目里，光加了这个参数，训练速度就提升了15%。代价是什么？锁页内存会占用更多系统内存，但这点代价完全值得。

3.3 prefetch_factor：提前多准备几份

这个参数很多人不知道。它控制每个worker提前预取多少个batch的数据。

默认值是2，意思是每个worker会提前准备2个batch的数据放在队列里。GPU用完一个，立刻就能拿到下一个。

# 默认预取2个batch
dataloader = DataLoader(dataset, batch_size=32, num_workers=4, 
                        pin_memory=True, prefetch_factor=2)

# 加大预取量，减少等待
dataloader = DataLoader(dataset, batch_size=32, num_workers=4, 
                        pin_memory=True, prefetch_factor=4)

为什么会这样？你想想看，每个batch的数据都占内存，预取越多，占的内存就越多。一般建议2到4之间，除非你的内存非常充裕。

3.4 完整配置示例

下面是我个人习惯的一套配置，适用于大多数场景：

import torch
from torch.utils.data import DataLoader

# 假设你的数据集
dataset = YourDataset()

# 优化后的DataLoader
dataloader = DataLoader(
    dataset,
    batch_size=32,
    shuffle=True,
    num_workers=4,          # CPU核心数的一半
    pin_memory=True,        # 开启锁页内存
    prefetch_factor=2,      # 预取2个batch
    persistent_workers=True # 训练完一个epoch不销毁worker进程
)

persistent_workers=True这个参数也值得提一下。它让worker进程在epoch结束后不销毁，省去了反复创建销毁的开销。如果你的数据集不大，epoch很多，这个参数能省不少时间。

3.5 知识体系总览

下面这张图帮你理清这三个参数的关系：

3.6 性能对比

我整理了一份实际测试数据，供你参考：

看到没？从默认到进阶优化，速度翻了一倍，GPU利用率从35%飙到85%。但过度优化反而会下降，因为进程切换和内存争抢开始拖后腿。

好了，数据加载这块就聊这么多。记住一句话：GPU是老板，数据加载是员工。员工干活慢，老板就得等。把这三个参数调好，让员工跑起来，老板才能满负荷工作。

想要更多资料，联系可可