4、数据加载与预处理：DataLoader原理、多进程加载、内存映射、数据增强的加速技巧、避免IO瓶颈

训练AI模型，说白了就是喂数据。但很多人把精力全花在调模型上，结果发现GPU利用率才30%，剩下的时间都在等数据。我刚开始做训练时也踩过这个坑——模型改了一百遍，速度就是上不去。后来才发现，瓶颈根本不在计算，而在IO。

这一章，咱们就把数据加载这块彻底聊透。

4.1 DataLoader 的工作原理

PyTorch 的 DataLoader，本质上就是个数据搬运工。它负责从硬盘上把数据读进来，做预处理，然后打包成 batch 送给 GPU。

核心流程其实就三步：

  1. 采样（Sampler）：决定从数据集中取哪些样本
  2. 读取与变换：把原始数据加载到内存，做预处理
  3. 组装成 batch：把多个样本拼成一个张量

我个人习惯把 DataLoader 理解成一个「生产者-消费者」模型。主进程是消费者，负责训练；子进程是生产者，负责准备数据。两者之间通过队列通信。

关键参数：

  • batch_size：每次取多少样本。不是越大越好，要看显存
  • shuffle：是否打乱顺序。训练集要打乱，验证集不用
  • num_workers：用几个子进程加载数据。这个最容易被忽略
  • pin_memory：锁页内存，能加速 CPU 到 GPU 的传输

我的经验：pin_memory=True 几乎总是有益的。尤其是在 GPU 训练时，它能减少数据传输的拷贝次数。我曾经在一个项目里忘了开这个，速度慢了将近20%。

4.2 多进程加载：num_workers 的正确用法

很多人以为 num_workers 设得越大越好。其实不是。你想想看，每个子进程都要独立加载数据，如果硬盘 IO 跟不上，开再多的进程也只是在排队等硬盘。

我建议的调参思路：

  • 先设成 0（单进程），跑一次看看 baseline
  • 然后逐步增加，观察 GPU 利用率的变化
  • 一般 4~8 个 worker 就够用了，再往上收益递减
  • 如果用了 SSD，可以适当多开；机械硬盘就别超过 4 个

注意：num_workers 开太多会导致内存暴涨。每个 worker 都会复制一份数据集对象。我曾经在 64GB 的机器上开了 32 个 worker，结果直接 OOM 了。嗯，血的教训。

还有一个容易被忽略的点：worker 的初始化开销。每个 worker 启动时都要加载一些库和资源。如果数据集很小，这个开销反而会拖慢速度。所以小数据集用单进程反而更快。

4.3 内存映射（Memory Mapping）

内存映射，说白了就是把硬盘上的文件直接映射到虚拟内存地址空间。这样读写文件就像读写内存一样，省去了系统调用的开销。

在 AI 训练中，最常用的场景就是加载大型数据集。比如 ImageNet，几百 GB 的图片，如果全部读到内存里，内存根本扛不住。但用内存映射，操作系统会按需加载，用多少读多少。

PyTorch 里怎么做？其实很简单：

import numpy as np
import torch

# 创建一个内存映射文件
data = np.memmap('data.bin', dtype='float32', mode='r', shape=(100000, 3, 224, 224))

# 直接当普通数组用
batch = torch.from_numpy(data[0:32])

我的建议：如果你的数据集超过内存的 50%，就用内存映射。它不会一次性占用所有内存，而是按需加载。我在做视频理解项目时，一个数据集就 200GB，全靠 memmap 撑下来的。

不过要注意，内存映射对随机访问特别友好，但顺序读取反而可能比普通文件慢。因为操作系统有预读机制，顺序读的时候会提前把后面的数据也加载进来。

4.4 数据增强的加速技巧

数据增强是防止过拟合的利器，但也是性能杀手。尤其是复杂的增强操作，比如随机裁剪、色彩抖动，计算量一点都不小。

我常用的加速技巧有这几个：

  1. 在 GPU 上做增强：如果 GPU 有空闲，可以把部分增强操作搬到 GPU 上。比如用 torchvision.transforms 的 GPU 版本
  2. 预计算增强：对于固定的增强策略，可以提前算好，存成文件。训练时直接读取
  3. 使用 DALI：NVIDIA 的 DALI 库专门为数据加载和增强做了优化，能充分利用 GPU
  4. 减少不必要的变换：有些增强操作对精度提升微乎其微，但计算量很大。比如高斯模糊，能不用就不用

一个实用的技巧：把数据增强分成「在线」和「离线」两部分。离线部分（比如缩放、裁剪）提前做好，存成统一尺寸。在线部分（比如翻转、色彩抖动）在训练时实时做。这样既保证了多样性，又不会太慢。

我记得有一次做图像分类，数据增强里有个随机旋转，每次都要插值计算，特别慢。后来我把旋转角度限制在 ±10 度，并且用最近邻插值代替双线性插值，速度提升了 3 倍，精度几乎没变。

4.5 避免 IO 瓶颈的实战策略

IO 瓶颈是训练加速的最大敌人。GPU 算得再快，数据送不进去也是白搭。我总结了一套排查和优化的流程：

第一步：诊断瓶颈

  • 看 GPU 利用率：如果低于 80%，说明数据加载跟不上
  • 看 CPU 利用率：如果 CPU 跑满，说明预处理太复杂
  • 看硬盘 IO：如果 IO 等待时间很长，说明硬盘是瓶颈

第二步：对症下药

症状 原因 解决方案
GPU 利用率低，CPU 空闲 硬盘太慢 换 SSD、用内存映射、预加载到内存
GPU 利用率低，CPU 跑满 预处理太重 简化增强、用 GPU 做预处理、增加 num_workers
GPU 利用率波动大 数据加载不均匀 使用 prefetch_factor、增大队列大小

第三步：终极方案——数据预加载

如果数据集能装进内存，那就一次性全部加载。这是最暴力的方法，也是最有效的。我做过一个实验：把 50GB 的数据集预加载到内存里，训练速度直接翻倍。

# 预加载到内存的示例
class PreloadedDataset(Dataset):
    def __init__(self, file_list):
        self.data = []
        for f in file_list:
            self.data.append(np.load(f))  # 一次性全部加载
    
    def __getitem__(self, idx):
        return self.data[idx]

警告：预加载虽然快，但要注意内存够不够。如果内存不够，操作系统会开始 swap，那速度反而更慢。我曾经在一个 32GB 的机器上预加载了 40GB 的数据，结果训练速度比原来还慢 5 倍。嗯，别学我。

4.6 知识体系总览

下面这张图总结了数据加载与预处理的核心逻辑：

数据加载与预处理核心流程 硬盘存储 原始数据 读取 内存映射 按需加载 预处理 数据增强 在线/离线 打包 DataLoader 多进程加载（num_workers） Worker 1 Worker 2 Worker 3 Worker 4 数据队列（Prefetch） pin_memory GPU 训练

这张图把整个流程串起来了。从硬盘到内存映射，再到数据增强，最后通过 DataLoader 的多进程机制送入 GPU。每一步都可能成为瓶颈，需要根据实际情况调整。

最后说一句：数据加载优化没有银弹。最好的方法就是 profiling。用 PyTorch 自带的 profiler 或者 nvidia-smi 看看瓶颈在哪，然后针对性地优化。我每次接手新项目，第一件事就是跑一次 profiling，看看数据加载到底占了多长时间。

想要更多资料，联系可可

← 上一章 📖 返回目录 下一章 →