4、数据加载与预处理：DataLoader优化、数据增强加速、内存与IO平衡
数据加载，说白了就是训练流程里的「送餐员」。
你模型再强，GPU再快，如果数据送不上来，那就是巧妇难为无米之炊。我见过太多团队，模型架构调得飞起，结果一跑训练，GPU利用率才30%，全在等数据。嗯，这章咱们就把这个瓶颈彻底打通。
4.1 DataLoader 的核心优化点
PyTorch 的 DataLoader 大家都会用，但用好的人不多。我个人习惯，上来先检查三个参数：num_workers、pin_memory、prefetch_factor。
核心原则：让数据加载和模型计算并行起来，互不等待。
4.1.1 num_workers 到底设多少？
这个参数决定了用几个子进程去加载数据。设少了，CPU 闲着；设多了，进程切换开销反而拖慢速度。
我在项目中踩过一个坑：直接设成 32，结果数据加载比单进程还慢。为什么？因为 CPU 核数只有 16，进程一多，都在抢资源，反而互相打架。
我的经验公式：
- CPU 密集型预处理（如图像解码、复杂变换）：
num_workers = CPU 核数 - IO 密集型（简单读取、少量变换）：
num_workers = 2~4 倍 CPU 核数 - GPU 训练时：建议从
num_workers = 4开始，逐步增加，观察 GPU 利用率
小技巧：用 nvidia-smi 看 GPU 利用率，如果低于 80%，先加 num_workers。如果加了没变化，说明瓶颈在别处。
4.1.2 pin_memory 和 prefetch_factor
pin_memory=True 这个选项，我建议你永远打开。它把数据固定在内存中，GPU 可以直接通过 DMA 读取，省去了 CPU 到 GPU 的拷贝时间。
至于 prefetch_factor，它控制每个 worker 预取多少个 batch。默认是 2，我一般调到 4 或 8。你想想看，如果每个 worker 提前准备好 8 个 batch，GPU 几乎永远不会饿着。
# 一个比较激进的配置，适合高性能场景
dataloader = DataLoader(
dataset,
batch_size=64,
shuffle=True,
num_workers=8,
pin_memory=True,
prefetch_factor=4,
persistent_workers=True # 避免每个 epoch 重新创建 worker
)注意：persistent_workers=True 会占用更多内存，但能省掉每个 epoch 启动 worker 的开销。如果你的内存够用，建议开启。
4.2 数据增强加速：别让预处理拖后腿
数据增强是提升模型泛化能力的好方法，但也是性能杀手。我曾经在一个图像项目里，用了 10 种随机增强，结果数据加载时间翻了 5 倍。说白了，增强太复杂，反而得不偿失。
4.2.1 用 GPU 做数据增强
既然 GPU 闲着也是闲着，不如让它分担一部分增强工作。PyTorch 的 torchvision.transforms 支持在 GPU 上执行部分操作。
import torch
from torchvision import transforms
# 在 GPU 上做随机裁剪和翻转
gpu_transform = transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
# 在训练循环中
images = images.cuda(non_blocking=True)
images = gpu_transform(images) # 在 GPU 上执行这样做的好处是：CPU 只负责读取和简单解码，复杂的几何变换交给 GPU。我实测过，对于 ResNet-50 训练，GPU 增强比 CPU 增强快了将近 30%。
4.2.2 使用 DALI 库
NVIDIA 的 DALI（Data Loading Library）是专门为加速数据加载和预处理设计的。它把整个 pipeline 都放在 GPU 上执行，包括图像解码、缩放、裁剪、颜色变换等。
from nvidia.dali.pipeline import Pipeline
import nvidia.dali.ops as ops
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 = ops.FileReader(file_root="./data")
self.decode = ops.ImageDecoder(device="mixed")
self.resize = ops.Resize(device="gpu", resize_x=224, resize_y=224)
self.cmn = ops.CropMirrorNormalize(
device="gpu",
output_dtype=types.FLOAT,
crop=(224, 224),
mean=[0.485 * 255, 0.456 * 255, 0.406 * 255],
std=[0.229 * 255, 0.224 * 255, 0.225 * 255]
)
def define_graph(self):
images, labels = self.input()
images = self.decode(images)
images = self.resize(images)
images = self.cmn(images)
return images, labels我的经验：DALI 适合大规模训练（比如 ImageNet 级别）。如果你的数据集只有几千张，用 PyTorch 原生 DataLoader 就够了，DALI 的配置成本反而划不来。
4.3 内存与 IO 平衡：别让数据撑爆内存
内存和 IO 是一对冤家。你想快，就得把数据多放内存里；但内存有限，放多了就 OOM。我见过有人把整个数据集加载到内存，结果 64GB 内存直接爆掉。
4.3.1 内存映射文件（mmap）
对于大型数据集（比如几十 GB 的图片），我建议用内存映射文件。它把磁盘文件映射到虚拟内存，操作系统按需加载，既快又省内存。
import numpy as np
import mmap
# 创建内存映射文件
data = np.memmap('large_dataset.bin', dtype='float32', mode='r', shape=(100000, 3, 224, 224))
# 使用时直接索引，操作系统自动处理缓存
batch = data[start:end]说白了，mmap 让操作系统帮你做缓存管理。你只管读数据，它自己决定哪些数据留在内存，哪些写回磁盘。
4.3.2 缓存策略：LRU 和预取
如果数据需要反复读取（比如每个 epoch 都要用），缓存是必须的。我常用的策略是：
- LRU 缓存：用
functools.lru_cache或自定义缓存，保留最近用过的数据 - 预取队列：用多线程提前加载下一个 batch 的数据
- 分层缓存：热数据放内存，冷数据放 SSD，归档数据放 HDD
from functools import lru_cache
@lru_cache(maxsize=1024)
def load_image(path):
# 从磁盘加载图片
return Image.open(path).convert('RGB')
# 训练时，重复的图片路径会命中缓存
for epoch in range(100):
for path in image_paths:
img = load_image(path) # 第二次调用时直接从缓存取避坑指南：我曾经用 LRU 缓存存了太多图片，结果内存不够，系统开始用 swap，训练速度直接掉到原来的十分之一。记住：缓存不是越大越好，要监控内存使用率。
4.4 知识体系总览
下面这张图，把数据加载与预处理的优化路径串起来了。你可以把它当作一个检查清单，遇到性能瓶颈时，按图索骥。
4.5 实战建议：从诊断到优化
如果你现在遇到数据加载慢的问题，我建议按这个顺序排查：
- 先看 GPU 利用率：低于 80%，说明数据加载是瓶颈
- 检查 num_workers：从 4 开始，逐步增加，直到 GPU 利用率不再提升
- 开启 pin_memory：这个几乎零成本，收益却很大
- 评估增强复杂度：如果增强太耗时，考虑移到 GPU 上做
- 监控内存使用：别让缓存吃掉所有内存，留 20% 给系统
一句话总结：数据加载优化的本质，就是让 CPU 和 GPU 各司其职，CPU 负责「喂」，GPU 负责「算」，中间用缓存和预取来平滑节奏。
好了，这章的内容就到这里。数据加载这块，说白了就是「细节决定成败」。你把这些参数调好了，训练速度翻倍不是梦。