文章目录
前言
在 PyTorch 的 DataLoader 中，num_workers 参数控制数据加载时的并行子进程数量。合理选择该参数可以显著提升数据加载效率，避免训练瓶颈。
1. 核心原则
CPU 核心数
num_workers 的理想值通常为 CPU 物理核心数的 2~4 倍。
例如：
若 CPU 有 8 核，建议设置为 4~8。
若 CPU 超线程（如 16 逻辑核心），建议设置为 8~16。
数据集大小
小数据集
小数据集（如内存可容纳）：设为 0（主进程加载）更高效，避免多进程开销。
大数据集
大数据集（需磁盘 I/O）：设为 4~8（根据 CPU 资源调整）。
预处理复杂度
若数据预处理（如数据增强）较复杂，适当增加 num_workers 可缓解计算压力。
2. 实验优化方法
监控 GPU 利用率
若 GPU 利用率低（如 <80%），可能是数据加载瓶颈，需增加 num_workers。
使用 nvidia-smi 或 PyTorch Profiler 观察 GPU 空闲时间。
逐步调参
从 num_workers=0 开始，逐步增加（如 2、4、8、16），记录每个 epoch 的耗时。
选择耗时最低且资源占用合理的值。
避免过度并行
若设置过高（如超过 CPU 核心数），可能导致进程切换开销增大，甚至内存溢出。
监控系统资源（如 htop 或 top），确保 CPU 和内存占用在安全范围内。
3. 具体场景建议
单GPU训练
=4~8 平衡并行加载与资源占用，适合大多数场景。
多GPU分布式训练
每个 GPU 2~4 总 num_workers = GPU 数量 × 单个 GPU 的推荐值，避免资源竞争。
轻量预处理
=2~4 数据加载简单（如仅读取图像），无需过高并行。
复杂预处理
=8~16 数据增强、特征提取等操作耗时，需更多子进程加速。
调试阶段
=0 避免多进程导致的调试问题（如断点失效、日志混乱）。
4. 注意事项
操作系统差异
Windows
Windows：多进程需将代码放在 if name == ‘main’: 中，否则可能报错。
Linux
Linux：支持更高效的多进程，可设置较高 num_workers。
内存限制
每个子进程会复制数据集到独立内存空间，若数据集过大，高 num_workers 可能导致 OOM。
共享文件系统
若数据存储在慢速磁盘或网络存储（如 HDD/NFS），增加 num_workers 可能收效甚微。
5. 示例代码
import multiprocessing
# 自动获取 CPU 核心数
cpu_cores = multiprocessing.cpu_count()
num_workers = min(4 * cpu_cores, 16) # 不超过 16
train_loader = DataLoader(
dataset=train_set,
batch_size=config["batch_size"],
shuffle=True,
num_workers=num_workers,
pin_memory=True # 启用锁页内存，加速 GPU 数据传输（需 GPU）
)
总结
默认推荐
默认推荐：从 num_workers=4 开始，逐步增加并观察训练速度。
关键指标
关键指标：确保 **GPU 利用率高（>90%）**且系统资源无瓶颈。
灵活调整
灵活调整：根据硬件、数据复杂度、预处理需求动态优化。