第4章数据加载与预处理迁移

第1章：数据加载与预处理迁移

数据加载，说白了就是整个训练流程的「水龙头」。水龙头开得太小，GPU 就得干等着。我见过太多团队，花大价钱买了 A100，结果 GPU 利用率不到 30%。为什么？数据加载卡脖子了。

这一章，我们就来聊聊怎么把这「水龙头」拧到最大。

从 CPU 多进程 DataLoader 说起

PyTorch 的 DataLoader，大家应该都用过。默认的 num_workers=0，数据在主进程里加载，GPU 吃一口等一口，效率极低。

后来我们学会了设置 num_workers>0。比如设成 4 或 8，让子进程提前把数据准备好。这确实有用，但有个坑——内存会暴涨。

我曾经踩过的坑：有一次我把 num_workers 设成 32，结果 256GB 内存直接爆了。每个 worker 都会复制一份数据，你想想看，32 个 worker 就是 32 份副本。

所以我的建议是：num_workers 一般设为 CPU 核心数的一半，最多不超过 2 倍。同时配合 prefetch_factor 参数控制预取数量。

# 一个比较稳妥的配置
DataLoader(
    dataset,
    batch_size=64,
    num_workers=8,        # 根据 CPU 核心数调整
    prefetch_factor=2,    # 每个 worker 预取 2 个 batch
    pin_memory=True,      # 锁页内存，加速 CPU→GPU 传输
)

这里 pin_memory=True 很关键。它把数据固定在物理内存里，GPU 可以直接 DMA 读取，省去了页面交换的开销。我个人习惯，只要用 GPU 训练，这个选项必开。

GPU Direct I/O：让数据直达显存

传统流程是这样的：

磁盘 → 系统内存（CPU 拷贝）
系统内存 → 显存（PCIe 传输）

两步走，中间有两次拷贝。GPU Direct I/O 能跳过第一步，让磁盘数据直接进显存。

为什么会这样？因为 GPU Direct 技术允许存储设备（比如 NVMe SSD）通过 PCIe 直接与 GPU 通信，绕过了 CPU 和系统内存。

  核心要点：GPU Direct I/O 不是银弹。它要求硬件支持（NVIDIA GPU + 支持 GPUDirect 的存储），而且只对大数据块传输有明显收益。小文件场景下，收益微乎其微。

我在项目中测试过，对于 4K 视频帧这种大文件，GPU Direct 能带来 20%-30% 的吞吐提升。但如果是 ImageNet 那种 256x256 的小图，反而因为初始化开销变慢了。

NVIDIA DALI：数据流水线的终极方案

DALI 是 NVIDIA 推出的数据加载库。它把整个预处理流水线搬到了 GPU 上。

你想想看，传统的 CPU 预处理：图像解码、随机裁剪、颜色抖动、归一化……这些操作全在 CPU 上跑。CPU 既要跑训练逻辑，又要处理数据，忙得不可开交。

DALI 的做法是：用 GPU 的 CUDA 核心来做这些操作。GPU 闲着也是闲着，不如让它干活。

import nvidia.dali as dali
from nvidia.dali.pipeline import Pipeline
import nvidia.dali.types as types

class SimplePipeline(Pipeline):
    def __init__(self, batch_size, num_threads, device_id):
        super().__init__(batch_size, num_threads, device_id)
        self.input = dali.ops.FileReader(file_root="./data")
        self.decode = dali.ops.ImageDecoder(device="mixed")  # 混合模式
        self.resize = dali.ops.Resize(device="gpu", resize_x=224, resize_y=224)
        self.normalize = dali.ops.Normalize(device="gpu",
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225])

    def define_graph(self):
        images, labels = self.input()
        images = self.decode(images)
        images = self.resize(images)
        images = self.normalize(images)
        return images, labels

注意这里的 device="mixed"。解码操作在 CPU 上做（因为 JPEG 解码对 GPU 不友好），但后续的 resize、normalize 全在 GPU 上。这就是 DALI 的聪明之处——混合流水线。

我的经验：DALI 的加速效果在 batch size 较大时最明显。如果你的 batch size 只有 16 或 32，CPU 预处理可能已经够用了。但当你用到 256 或 512 的 batch，DALI 的优势就出来了——GPU 利用率能从 40% 飙升到 90% 以上。

内存映射文件与零拷贝技术

最后聊聊内存映射（mmap）。这技术其实不新，但在深度学习场景下特别好用。

传统读取文件：先 open，然后 read 到内存 buffer，再解析。每次 read 都是一次系统调用，开销不小。

mmap 的做法是：直接把文件映射到进程的虚拟地址空间。你访问文件就像访问内存一样，缺页时操作系统自动加载。

import mmap
import numpy as np

# 假设有一个 10GB 的二进制文件，存储了 float32 数据
with open("large_data.bin", "rb") as f:
    # 内存映射，不实际加载到内存
    mmapped_file = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ)
    # 直接当作 numpy 数组访问
    data = np.frombuffer(mmapped_file, dtype=np.float32)
    # 使用 data 进行训练...
    # 操作系统会自动按需加载页面

这样做的好处：

节省内存：只加载实际访问的部分
零拷贝：数据从磁盘到进程空间，没有中间 buffer
共享内存：多个进程可以映射同一个文件，共享物理内存

注意：mmap 不是万能的。如果你的数据访问模式是随机的、跳跃的，mmap 会导致大量缺页中断，反而比 read 还慢。它最适合顺序访问的场景——比如训练时按顺序读取样本。

嗯，这里要补充一点。零拷贝技术不止 mmap 一种。还有 sendfile、splice 等系统调用，它们在内核态直接传输数据，不经过用户空间。但在深度学习场景下，mmap 是最实用的。

知识体系总览

下面这张图总结了本章的核心逻辑：

从这张图可以看出，三种方案不是互斥的。我个人的建议是：

小规模实验（单卡、小 batch）：CPU DataLoader + pin_memory 就够了
中等规模（多卡、中等 batch）：上 DALI，把预处理搬到 GPU
大规模生产（集群、大 batch）：DALI + mmap + GPU Direct 组合拳

记住一个原则：不要让 GPU 等数据。数据流水线的吞吐量，至少要达到模型推理速度的 1.5 倍。否则 GPU 就会频繁进入空闲状态，利用率上不去。

好了，这一章的内容就到这里。数据加载是性能优化的第一道关卡，搞定了它，后面的路就好走了。

📚 课程目录

第1章：数据加载与预处理迁移

从 CPU 多进程 DataLoader 说起

GPU Direct I/O：让数据直达显存

NVIDIA DALI：数据流水线的终极方案

内存映射文件与零拷贝技术

知识体系总览