训练深度学习模型的提速方法

jackylzh2026-01-04 13:44

复制代码

1. nw = min([os.cpu_count() // max(nd, 1), batch_size if batch_size > 1 else 0, workers])

在PyTorch数据加载配置中，nw 通常代表数据加载器（DataLoader）的工作进程数（num_workers），用于控制数据预取的并行度。表达式 nw = min([os.cpu_count() // max(nd, 1), batch_size if batch_size > 1 else 0, ...]) 是一个常见的动态设置逻辑，旨在根据系统资源和训练配置自动优化 num_workers 的值。‌

‌该表达式通过最小值函数综合多个约束条件： ‌ 首先，os.cpu_count() // max(nd, 1) 计算可用逻辑CPU核心数，其中 nd 通常表示分布式训练的设备数（如GPU数量），若 nd 为0则默认使用全部核心，这避免了多进程数据加载与模型训练争夺CPU资源；其次，batch_size if batch_size > 1 else 0 将 num_workers 限制为 batch_size（当 batch_size 大于1时），因为小批量（如 batch_size=1）时多进程可能引入额外开销。实际代码中常包含更多约束，例如限制 num_workers 不超过某个值（如8）或根据显存容量调整，以平衡数据加载速度与系统稳定性。

修改训练命令中的 workers 参数：可以为CPU核心数的1.5倍，对应数据加载进程数。

例如：python train.py --data coco.yaml --epochs 300 --batch-size 64 --workers 12

tils/dataloaders.py的create_dataloader函数（第160行），创建数据加载器时设置num_workers。

默认为 nw = min( $os.cpu_count() // max(nd, 1), batch_size if batch_size \> 1 else 0, workers$ ) # 自动计算worker数

并确保pin_memory=True（第67行全局设置），可减少CPU到GPU的数据传输延迟。

4. 自动批处理大小

通过--batch-size -1让系统自动计算最佳批处理大小。

例如：python train.py --data coco.yaml --epochs 300 --batch-size -1

原理是：

原理是train.py第243行调用的check_train_batch_size函数，会通过内存测试找到最大可行批大小：

if RANK == -1 and batch_size == -1: # single-GPU only

batch_size = check_train_batch_size(model, imgsz, amp)

使用 NVIDIA-SMI 命令，实时查看训练程序时的gpu占用率

nvidia-smi dmon

从任务管理器，查看CPU利用率。

参考：超实用！YOLOv5 GPU利用率提升指南：从50%到95%的实战技巧-CSDN博客

https://blog.csdn.net/gitblog_00918/article/details/151241012