3. 数据加载参数配置（一）：batch_size、num_workers、shuffle、pin_memory详解
数据加载，这步看着不起眼，但很多时候训练慢、收敛不稳，根儿就在这几个参数上。我见过不少同学，模型搭得挺漂亮，结果数据加载成了瓶颈，GPU利用率上不去，空转浪费钱。
这节我把四个最核心的参数掰开揉碎了讲。batch_size、num_workers、shuffle、pin_memory，每个都有它的脾气和最佳实践。
3.1 batch_size：给GPU“一口吃多少”
batch_size就是个最直接的参数——每次训练迭代往GPU送多少个样本。太小，梯度震荡大；太大，显存爆了。
核心影响有三个：
- 收敛稳定性：batch太小，梯度的方差大，loss曲线像过山车。
- 显存占用：基本是线性关系。batch 64比batch 32多占一倍显存。
- 训练速度：不是越大越快，太大反而可能因为梯度更新次数变少而收敛变慢。
我个人习惯，先按显存上限试探。比如一块24G的RTX 3090，混精度训练ResNet-50，batch 256差不多就到顶了。再往上就要boom。
| batch_size范围 | 典型场景 | 优缺点 |
|---|---|---|
| 1~8 | 超大模型（LLM、3D视觉）、在线学习 | 梯度噪声大，但省显存，可做per-sample处理 |
| 16~64 | 常见CV模型（ResNet、ViT） | 比较均衡，收敛稳定性和速度都可接受 |
| 128~512 | 多GPU训练、大batch收敛优化 | 梯度估计更准，但要小心sharp minima，可能需要调学习率 |
| 1024+ | 大规模分布式训练（Google常用） | 必须配合learning rate warmup和cosine decay，否则直接发散 |
为什么大batch需要调整学习率？梯度是从更多样本平均出来的，方差小了，但每一步的方向可能更“滑”。说白了，需要更大的步长才能有效探索。
# PyTorch DataLoader配置示例
from torch.utils.data import DataLoader
train_loader = DataLoader(
dataset=train_dataset,
batch_size=64, # 每次64张图
shuffle=True,
num_workers=4,
pin_memory=True
)3.2 num_workers：CPU帮GPU“提前备料”
这个参数控制的是数据加载的并行进程数。每个worker独立地去读盘、解压、做数据增强，然后往一个共享队列里放。
我之前遇到过一个项目，GPU利用率才30%上下，怎么看都是瓶颈在数据侧。一查，num_workers=0，所有加载都在主进程串行做，CPU等IO，GPU等人。改成4之后，利用率直接拉到85%。
关于选几个worker，我的经验公式：
num_workers = 2 * CPU核心数 # 有超线程的情况
num_workers = 1 * CPU核心数 # 无超线程或IO密集场景
num_workers = 0 # debug或小数据集- 设到CPU核心数两倍时，通常读盘和计算都能掩盖掉延迟。
- 设太大（比如32核机子上设64个worker），反而因为上下文切换和内存争抢导致变慢。
- 设置成0会让所有加载跑在主进程，调试方便，但训练效率很低。
你想想看，每个worker都是一个完整的Python进程，它们复制了一份数据和transform链。如果transform很重（比如RandAugment加mixup），进程数就得适当减少，不然CPU要被算爆。
nvidia-smi看GPU利用率，再用htop看CPU和IO状态。如果Volatile GPU-Util经常跳变，就从调大num_workers开始。
3.3 shuffle：别让模型“偷懒”
shuffle控制是否在每个epoch开始时打乱数据集。这个参数虽小，但影响很大。
训练集：必须shuffle=True。如果不打乱，模型会学到样本的顺序信息。举个例子，数据集按类别排好序——前500张都是猫，后500张都是狗。模型前500步一直在学猫，梯度往一个方向带偏，等看到狗的时候前面学的已经固化了。
验证集：必须shuffle=False。原因就一条——结果要可复现。如果每次验证顺序不同，指标之间的对比就没法做。你没法区分是模型变好了还是运气变好了。
train_loader = DataLoader(
train_set,
batch_size=32,
shuffle=True, # 训练时打乱
num_workers=4
)
val_loader = DataLoader(
val_set,
batch_size=32,
shuffle=False, # 验证时固定顺序
num_workers=4
)3.4 pin_memory：给数据传输“开高速”
pin_memory=True表示把数据加载到锁页内存（pinned memory）里。普通内存可能会被操作系统交换到磁盘（swap），但锁页内存会固定在物理内存中不走。
说白了，GPU显存和CPU内存之间有一条DMA通道。DMA只能访问锁页内存。如果你开pin_memory=True，数据传输就可以走硬件DMA，速度比CPU搬数据快一个数量级。
我之前用过一个小技巧：在数据加载前把一块小区域pin住，配合num_workers，把预处理后的tensor直接放进去，传输到GPU几乎零拷贝。
# 标准用法
train_loader = DataLoader(
...
pin_memory=True, # 传输更快
pin_memory_device='cuda:0' # PyTorch 2.0+ 可指定目标设备
)3.5 这四个参数怎么一起调？
实践下来，我有一个从零开始调优的套路，贴出来给你参考一下。
| 步骤 | 操作 | 目标 |
|---|---|---|
| 1 | 设batch_size为显存能容纳的最大值（试探法） | 最大化吞吐 |
| 2 | 设num_workers=0，跑一个epoch | 建立baseline，验证代码正确 |
| 3 | 逐步调大num_workers（2, 4, 8...），看GPU利用率 | 找到CPU和IO的平衡点 |
| 4 | 开启pin_memory=True | 加速CPU→GPU传输 |
| 5 | 调整prefetch_factor（默认2），减少worker空闲 | 进一步消除IO等待 |
# 一个生产中常用的组合
loader = DataLoader(
dataset,
batch_size=128,
shuffle=True,
num_workers=8,
pin_memory=True,
drop_last=True,
prefetch_factor=2, # 每个worker提前加载2个batch
persistent_workers=True # epoch结束时worker不销毁，减少创建开销
)