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

做量化模型训练，我最怕看到什么？GPU利用率只有20%，显存占用倒是挺高，但计算单元全在摸鱼。你想想看，一张A100一小时几十块钱，结果一半时间在等数据加载——这钱花得冤不冤？

数据加载优化，说白了就是解决「GPU等饭」的问题。我见过太多团队，模型写得花里胡哨，结果数据管道一塌糊涂。今天咱们就把DataLoader的每个参数掰开揉碎，看看怎么让GPU吃得饱、跑得快。

DataLoader工作原理：它到底在干嘛？

PyTorch的DataLoader，本质上是一个数据搬运工。它从硬盘读取数据，做预处理，然后送到GPU。但这里有个关键问题：数据读取是IO密集型，模型训练是计算密集型。这两件事天然不匹配。

我习惯把DataLoader的工作流程拆成三步：

1. 采样（Sampler）：决定取哪些样本
2. 加载（Loader）：从硬盘读到内存
3. 预处理（Transform）：归一化、裁剪、转Tensor

这三步是串行的。如果每一步都等上一步完成，GPU就只能干瞪眼。所以PyTorch引入了多进程机制——让多个worker并行干活。

举个例子，假设你有一个worker，它一次只能搬一箱数据。搬完、处理完、送到GPU，然后才能搬下一箱。这中间的空闲时间，GPU就在等。如果有4个worker，它们可以同时搬4箱，GPU几乎不用等。

num_workers调优：到底开多少个合适？

这是最常被问的问题。num_workers设多少？8？16？32？

我的回答是：不是越多越好。每个worker都是一个独立进程，它们之间需要通信、需要同步。worker太多，反而会因为进程切换和内存争抢拖慢速度。

我有个简单的经验公式：

num_workers = 4 × GPU数量

但这只是起点。实际调优时，我会这样做：

1. 从2开始，逐步增加
2. 观察GPU利用率的变化
3. 当GPU利用率不再提升时，就是最佳值

我曾经在一个项目里，把num_workers从4调到8，训练速度提升了40%。但调到16时，反而慢了10%。为什么？因为CPU成了瓶颈，进程切换开销超过了并行收益。

还有一个坑：Windows上num_workers设0最稳。因为Windows的进程创建机制和Linux不同，多进程容易出问题。我建议在Linux上训练，省心很多。

pin_memory使用：锁页内存的妙用

pin_memory=True，这个参数很多人知道要设，但不知道为什么。我来解释一下。

CPU内存分两种：

• 可分页内存（Pageable）：默认情况，数据可能被操作系统换到硬盘上
• 锁页内存（Pinned）：数据固定在物理内存中，不会被换出

GPU从CPU拷贝数据时，如果数据在可分页内存里，需要先拷贝到锁页内存，再拷贝到GPU。多了一步操作。如果数据已经在锁页内存里，直接拷贝到GPU，速度更快。

说白了，pin_memory=True就是省掉中间商赚差价。

# 推荐做法
dataloader = DataLoader(
    dataset,
    batch_size=64,
    shuffle=True,
    num_workers=4,
    pin_memory=True  # 一定要开
)

我测试过，开启pin_memory后，数据从CPU到GPU的传输速度能提升2-3倍。对于大batch size的训练，效果尤其明显。

prefetch_factor设置：预取多少合适？

prefetch_factor是PyTorch 1.7之后引入的参数。它控制每个worker预取多少个batch的数据。

默认值是2。什么意思？每个worker会提前加载2个batch的数据，放在缓冲区里。GPU用完了当前batch，直接拿下一个，不用等。

这个参数怎么调？我的经验是：

• 数据加载慢（比如从网络读取）：prefetch_factor设大一点，比如4或8
• 数据加载快（比如从SSD读取）：默认的2就够了
• 内存紧张：设小一点，甚至设1

我曾经在一个分布式训练项目里，把prefetch_factor从2调到8，训练速度提升了15%。但代价是内存占用翻了一倍。嗯，这里要注意：预取越多，内存占用越大。

# 数据加载慢的场景
dataloader = DataLoader(
    dataset,
    batch_size=64,
    num_workers=4,
    pin_memory=True,
    prefetch_factor=4  # 每个worker预取4个batch
)

内存映射文件：终极IO优化方案

前面说的都是小打小闹。如果你真的想榨干IO性能，内存映射文件（Memory-Mapped File）是终极方案。

原理很简单：把文件直接映射到进程的虚拟地址空间。读写文件就像读写内存一样，不需要显式的read/write系统调用。

在量化交易场景下，这个技术特别有用。因为我们的数据通常是结构化、固定大小的（比如每天的K线数据）。用内存映射文件，可以实现近乎零拷贝的数据加载。

import numpy as np
import torch
from torch.utils.data import Dataset

class MMapDataset(Dataset):
    def __init__(self, mmap_path, shape, dtype=np.float32):
        # 创建内存映射文件
        self.data = np.memmap(
            mmap_path,
            dtype=dtype,
            mode='r',
            shape=shape
        )
    
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        # 直接读取，没有系统调用开销
        sample = torch.from_numpy(self.data[idx].copy())
        return sample

我做过对比测试：

加载方式	100万条数据加载时间	内存占用
普通文件读取	12.3秒	高（全部加载到内存）
内存映射文件	0.8秒	低（按需加载）
内存映射+预取	0.3秒	中（部分缓存）

看到差距了吗？内存映射文件比普通读取快了15倍。而且内存占用还更低——因为操作系统只把真正用到的部分加载到内存。

实战组合拳：一个完整的优化方案

说了这么多，咱们来一个完整的配置示例：

from torch.utils.data import DataLoader

# 优化后的DataLoader配置
dataloader = DataLoader(
    dataset=mmap_dataset,      # 使用内存映射数据集
    batch_size=256,
    shuffle=True,
    num_workers=4,             # 根据CPU核心数调整
    pin_memory=True,           # 开启锁页内存
    prefetch_factor=4,         # 预取4个batch
    persistent_workers=True,   # worker进程不销毁，复用
    timeout=60                 # 超时保护
)

这套配置我用了很久，效果很稳定。GPU利用率从原来的40%提升到了85%以上。

最后说一句：数据加载优化没有银弹。不同的硬件、不同的数据格式、不同的模型，最优配置都不一样。我的建议是：先理解原理，再动手测试，最后找到最适合你的方案。

好了，这一章的内容就到这里。记住，GPU很贵，别让它闲着等饭。

---