3. 数据加载优化

聊到训练吞吐，很多人第一反应就是调大 batch size、换更快的 GPU。但我得说，很多时候瓶颈根本不在计算上——数据加载才是那个拖后腿的隐形杀手。

我自己就踩过这个坑。有一次跑一个图像模型，GPU 利用率死活上不去，一直在 30% 左右晃悠。我查了半天，最后发现 DataLoader 的 num_workers 设成了默认值 0。说白了，数据加载完全是在主进程里串行跑的，GPU 大部分时间都在干等。改完参数之后，利用率直接飙到 90% 以上。

所以这一章，咱们就好好聊聊 DataLoader 的那些门道。

DataLoader 到底怎么工作的？

先看个最简单的例子：

from torch.utils.data import DataLoader, Dataset

class MyDataset(Dataset):
    def __init__(self, data):
        self.data = data
    
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        return self.data[idx]

dataset = MyDataset(list(range(100)))
dataloader = DataLoader(dataset, batch_size=8, shuffle=True)

这段代码大家应该都很熟悉。但 DataLoader 内部是怎么把数据喂给模型的？我简单拆解一下：

• 采样器（Sampler）：决定从数据集中取哪些样本。shuffle=True 时用的是 RandomSampler，否则是 SequentialSampler。
• 批处理（Batch Sampler）：把采样器输出的索引组合成一个个 batch。比如 batch_size=8，那就每 8 个索引一组。
• 数据加载（Dataset.__getitem__）：根据索引从磁盘或内存中读取实际数据。
• 整理（collate_fn）：把 batch 里的样本拼成一个张量。默认的 collate_fn 就是简单的堆叠。

整个过程看起来挺直白的，对吧？但性能问题就出在「数据加载」这一步——如果每次都要从磁盘读文件，那速度可就慢得离谱了。

num_workers 调优：多进程到底开几个？

num_workers 控制的是数据加载的子进程数量。默认是 0，表示在主进程中加载数据。我建议至少设成 1，除非你的数据集小到可以忽略加载时间。

那到底设多少合适？我个人的经验是：

硬件配置	推荐 num_workers	说明
单 GPU，CPU 核数少（4-8）	2-4	别开太多，否则进程切换开销反而更大
单 GPU，CPU 核数多（16+）	4-8	我一般从 4 开始试，看 GPU 利用率
多 GPU，CPU 核数多	每个 GPU 对应 2-4	注意总进程数别超过 CPU 核数太多

调优方法其实很简单：先设成 0 跑一次，记录 GPU 利用率。然后逐步增加 num_workers，每次加 2，观察 GPU 利用率的变化。当利用率不再明显提升时，那就是最佳值了。

prefetch_factor：提前加载多少数据？

prefetch_factor 这个参数很多人容易忽略。它控制的是每个 worker 提前预加载多少个 batch 的数据。默认是 2，意味着每个 worker 会提前准备 2 个 batch。

为什么要预加载？说白了就是「以空间换时间」。GPU 在算当前 batch 的时候，worker 已经在准备下一个 batch 了。这样 GPU 算完一个，下一个马上就能接上，不用干等。

# 默认 prefetch_factor=2
dataloader = DataLoader(dataset, batch_size=32, 
                        num_workers=4, prefetch_factor=2)

# 如果内存够，可以试试更大的值
dataloader = DataLoader(dataset, batch_size=32,
                        num_workers=4, prefetch_factor=4)

我记得有一次在 32GB 内存的机器上训练，batch_size 设了 128，图片每张 500KB 左右。prefetch_factor 设成 8 之后，内存直接爆了。嗯，后来我老老实实改回了 2。

内存映射文件：让数据加载飞起来

内存映射（Memory Mapping）是我个人非常喜欢的一个技巧。它把磁盘上的文件直接映射到虚拟内存地址空间，读写的时候就像操作内存一样，不需要显式的 read/write 调用。

在 PyTorch 里，可以用 torch.from_numpy 配合 np.memmap 来实现：

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):
        # 直接从内存映射中读取，不需要磁盘 I/O
        return torch.from_numpy(self.data[idx].copy())

# 使用示例
dataset = MMapDataset('data.bin', (10000, 3, 224, 224))
dataloader = DataLoader(dataset, batch_size=32, num_workers=4)

但要注意，内存映射文件也有坑。我遇到过一个问题：如果数据集太大，超过了虚拟地址空间（32 位系统下只有 4GB），就会报错。解决方案是用 64 位 Python，或者把数据集分片。

综合调优策略

说了这么多，咱们总结一下实际调优的步骤：

1. 先确认瓶颈：用 nvidia-smi 看 GPU 利用率，如果低于 80%，大概率是数据加载的问题。
2. 调整 num_workers：从 2 开始，逐步增加，直到 GPU 利用率不再提升。
3. 调整 prefetch_factor：在内存允许的情况下，适当增大。
4. 考虑内存映射：如果数据集是固定大小的（比如图像、特征向量），强烈建议用 mmap。
5. 检查 collate_fn：如果自定义了 collate_fn，确保它没有成为新的瓶颈。

下面这张图展示了数据加载的完整流程和优化点：

最后说一句，调优不是一锤子买卖。换了数据集、改了模型、甚至换了机器，都可能需要重新调整这些参数。我自己的习惯是每次训练前先跑一个 mini 测试，确认数据加载不是瓶颈，再开始正式训练。

嗯，数据加载这块就聊到这儿。记住一句话：别让 GPU 闲着，也别让 CPU 累死。找到那个平衡点，你的训练吞吐就能上一个台阶。