★ FEATURED ARTICLE 前阵子有同学跑YOLOv5训练目标检测模型五天五夜没关机第二天就跑过来问我为什么这么慢是不是显卡不行我看了眼他的启动命令单卡RTX 3080数据是小批量图像按道理不该这么难受。最后发现症结根本不在显卡上而是数据加载线程没调、混合精度没开加上每两个epoch就全量验证一次。这篇文章我把这两年调YOLOv5训练速度的思路完整盘一遍从硬件环境、数据加载、模型配置到多卡分布式按加速收益从高到低排列。你不需要照着全部改先定位自己卡在哪一环再去动对应的地方效果会立竿见影。1. 训练慢先别急着加卡瓶颈往往在你看不见的地方1.1 训练速度慢的五个典型表现和成因很多人的第一反应是“显卡太弱”于是急着上多卡、换A100。但如果GPU根本没有被喂饱加再多卡也是浪费钱。我习惯先看几个现象再判断到底是哪一类问题。训练慢的常见表现大致分五类GPU利用率低50%甚至更低CPU反而跑满。这是典型的数据加载瓶颈GPU大部分时间在等CPU把图像处理好送过来。GPU利用率很高但每秒迭代数就是上不去。这种情况一般是模型计算量太大或者batch size太小没法吃满显存带宽。显存不够只能被迫调小batch size。迭代次数变多总训练时间呈线性增长。训练过程每隔一段时间卡一下。可能是验证集评估、日志写入、checkpoint保存或磁盘IO波动看起来像“呼吸停顿”。前几个epoch很慢后面就正常了。通常是首次缓存图像、anchors计算或模型预热带来的额外开销。表现优先排查方向处理建议GPU利用率低、CPU满载数据加载线程不足调大--workers开启缓存GPU利用率高但速度依然慢batch size太小调大batch或换更小模型显存不足batch调不大输入尺寸/模型太大降--img换小模型开AMP训练周期卡顿验证/日志/checkpoint降低验证频率关闭在线日志首个epoch特别慢缓存和预处理使用--cache ram或disk拿着一张症状表去对比自己的机器很快就能锁定方向。而不是一上来就把所有参数都改了那样反而不知道是哪一步起的作用。1.2 动手前的基准测试把“慢”量化成数字人脑对“慢”的感知是模糊的所以动手优化前我会先跑一个基准测试记录几个硬指标每个epoch的耗时每秒钟处理的迭代次数it/sGPU利用率和显存占用数据加载阶段是否出现明显等待YOLOv5训练时会在终端打印it/s和batches/s这是最直接的参考。但只看这个还不够我习惯另外开一个终端跑监控命令watch -n 1 nvidia-smi这个命令每秒刷新一次GPU利用率、显存、温度、功耗。如果训练过程中GPU利用率经常掉到80%以下那就说明存在等待。进一步定位可以看CPU占用top -o %CPU或者用nvidia-smi dmon输出连续的监控数据nvidia-smi dmon -f 60基准测试还有一个好处能区分“训练本身慢”和“流程整体慢”。我遇到过有人训练非常快但每次启动都要等十几分钟读数据因为数据集放在机械硬盘里没有做缓存。这种情况下优化训练参数没有任何意义真正要解决的是IO路径。做完基准测试把问题归类再往下针对性加速。2. 混合精度与运行环境几分钟就能见效果的加速项2.1 混合精度训练AMP的正确打开方式混合精度是YOLOv5训练加速里投入产出比最高的一个开关。简单说就是把一部分计算从FP32精度降到FP16显存占用几乎减半计算速度也能明显提升尤其是消费级显卡效果格外明显。YOLOv5官方训练脚本已经内置AMP支持。在较新的版本中默认会自动检测是否开启但如果你的版本较旧或者保险起见手动加上参数即可python train.py --img 640 --batch 16 --epochs 100 --data data.yaml --weights yolov5s.pt --amp我在一张总显存8GB的单卡上测过开了AMP之后it/s从2.3提升到4.1显存占用从6.8GB降到4.2GB。训练时间缩短了快一半精度损失几乎可以忽略。这里有几个容易踩的坑开启AMP后loss出现nan先别怀疑混合精度本身大概率是学习率太高或训练过程动了数据分布。可以先把学习率降低一个数量级或者配合warm-up训练前3个epoch。如果你的显卡太老比如一些不完整支持FP16加速的卡AMP未必有收益甚至可能更慢。可以关掉AMP对比一下。AMP的loss scaling机制一般会自动调节不需要手动干预。但偶尔会遇到loss为nan后整体卡死的情况同步检查驱动版本和CUDA环境会更有帮助。2.2 CUDA环境与驱动的匹配检查清单环境问题导致的训练慢排查起来最隐蔽。因为错误通常不会报出来只是悄悄用CPU在跑或者显卡驱动和PyTorch没有对齐。第一步确认PyTorch真的能使用GPUpython -c import torch;print(torch.__version__);print(torch.cuda.is_available());print(torch.cuda.get_device_name(0))如果cuda.is_available()返回False那训练会全部落到CPU上慢到怀疑人生。这种情况多是因为安装PyTorch时用的CUDA版本和系统驱动不匹配。PyTorch版本常用CUDA版本建议驱动版本区间1.13 / 2.0CUDA 11.7 / 11.85152.1 / 2.2CUDA 12.15302.3CUDA 12.1 / 12.4550安装PyTorch时用官方网站给出的命令安装不要自己手动折腾系统CUDA路径这是最省事的方法。我见过太多次因为手动改了系统CUDA符号链接结果训练时兼容性出问题。另外多卡机器上建议用环境变量控制可见显卡export CUDA_VISIBLE_DEVICES0,1,2,3这样能避免程序经常扫描所有显卡导致的初始化开销也让显存分配更可控。3. 数据加载与预处理GPU在等CPU这是最常见的隐性瓶颈3.1 num_workers与pin_memory怎么调才合理YOLOv5训练时数据增强和图片预处理都在CPU上完成GPU则在等CPU投喂数据。--workers参数控制的是用于数据加载的子进程数量这个参数设置不合理会直接造成GPU利用率忽高忽低。YOLOv5默认的--workers 8在绝大多数机器上都能用但并不是越大越好。子进程太多会导致操作系统频繁切换进程反而增加开销。我的建议是物理核数在8核以内的机器设成4或816核及以上的机器可以先从8试起观察GPU利用率再决定要不要加到12或16Windows系统下num_workers不要设太高多进程机制带来的开销比Linux下更明显。开启pin_memory默认开启可以把数据从页锁定内存直接传输到GPU减少一次拷贝。如果内存充裕不用管它如果内存特别紧张可以加上--no-pin-memory关闭。我自己常用的一组命令python train.py --data data.yaml --weights yolov5s.pt --img 640 --batch 16 --workers 8 --cache ram要注意的是--workers最优值受数据增强复杂度影响。比如YOLOv5默认开Mosaic增强Mosaic合并了四张图对CPU算力要求比较高。CPU偏弱的机器上即使num_workers拉高GPU利用率也未必能上去。3.2 缓存策略与图像尺寸对训练速度的影响数据读取这条链路很多人会忽略文件系统本身。如果数据集放在机械硬盘每次训练都要随机读取几万张图片速度会被磁盘的随机IO直接拖垮。YOLOv5提供了--cache选项可以把图像一次性加载到内存或SSD缓存中训练期间不再从原始路径逐张读取python train.py --data data.yaml --weights yolov5s.pt --img 640 --batch 16 --cache ram--cache ram适合中小规模数据集数据一次性读进内存训练期间完全不再读盘。首次加载会比较慢但之后训练过程的等待时间几乎清零。如果内存不够大可以用--cache disk效果不如内存但也能减少机械硬盘的随机IO。图像尺寸--img参数对训练速度的影响同样很大。默认640px如果跑VOC这类公共数据集640是一个平衡点。如果数据集里目标普遍偏大、分辨率要求不高可以先降到512甚至416跑通整个流程最后再用更高分辨率精调几个epoch。这里要说个反向经验不要盲目把尺寸调大追求精度。有一次我把--img从640升到1280训练时间直接变成原来的4倍而最终mAP只提升了不到1个点。数据增强手段已经能带来足够的多样性高分辨率只在最后微调阶段值得上。4. 模型配置与训练策略换个思路也能省下大把时间4.1 模型尺寸选择不是越大越好YOLOv5提供n/s/m/l/x五个尺寸版本很多人习惯“直接用最大的”结果就是训练时间翻倍、显存告急。实际上模型尺寸应该跟数据集规模和任务难度匹配。模型参数量训练耗时相对值适用场景YOLOv5n约1.9M0.5x数据量小、对实时性要求高YOLOv5s约7.2M1x大多数目标检测任务的起点YOLOv5m约21.2M约1.6x数据量中等、精度要求较高YOLOv5l约46.5M约2.5x大数据集、追求更高精度YOLOv5x约86.7M约4x大规模数据、计算资源充足我做过一个对比实验在几百张图片的小数据集上用v5x训练相比v5s耗时多出将近3倍精度只提高了不到2个百分点。训练数据量有限时大模型的优势根本发挥不出来反而更容易过拟合。实操建议是新项目先用v5s把整个流程跑通确认数据、标签、验证逻辑都没问题再根据训练曲线决定是否升级到v5m。这样最省时间也最稳妥。4.2 冻结主干与分阶段训练迁移学习省时间YOLOv5预训练权重是在大规模数据上训练出来的主干部分的特征提取能力已经很强。从头训练一个检测模型前期大部分计算量其实花在重复学习边缘、纹理这些基础特征上完全没必要。--freeze参数可以冻结主干部分只训练后面的检测头。冻结的层不参与反向传播计算量大幅下降显存占用也会变小python train.py --data data.yaml --weights yolov5s.pt --img 640 --batch 16 --epochs 50 --freeze 10--freeze 10表示冻结模型前10层一般就是主干网络backbone。对于迁移学习场景可以先冻结主干快速训练10个epoch让检测头适配你的数据集之后再解冻主干用较低学习率微调整个模型python train.py --data data.yaml --weights runs/train/exp/weights/last.pt --img 640 --batch 16 --epochs 20 --freeze 0 --lr 0.0001分阶段训练时要注意冻结主干会同时冻结相关BatchNorm层的统计量更新如果你的数据集跟预训练数据的风格差异很大第一阶段不宜太长否则后面解冻时模型波动会比较大。还有一个实用习惯正式跑长训练之前先跑2~3个epoch确认pipeline顺畅再启动完整训练。这样能提前发现问题避免在错误配置上白跑几十个小时。5. 多卡分布式训练从单卡到多卡的关键调整5.1 DDP启动参数与显存分配单卡训练到瓶颈后多卡分布式训练是自然的下一步。YOLOv5原生支持PyTorch DDPDistributed Data Parallel和老的DataParallel比起来DDP每个进程持有独立的模型和优化器只在梯度阶段同步效率高得多。单机四卡启动命令如下CUDA_VISIBLE_DEVICES0,1,2,3 python -m torch.distributed.run --nproc_per_node4 train.py --data data.yaml --weights yolov5s.pt --img 640 --batch 64更老版的写法是python -m torch.distributed.launch --nproc_per_node4 train.py --data data.yaml --weights yolov5s.pt --img 640 --batch 64两种写法本质上是一样的新版推荐用torch.distributed.run。需要特别注意YOLOv5的--batch-size在DDP模式下指的是“所有卡上的总batch”。四卡跑--batch 64平均每张卡分到16张图。如果你的显存只够每张卡跑8张那就把总batch调小比如--batch 32。我遇到过几次DDP启动后卡死的问题多半是端口冲突或残留进程占用资源。可以换个master端口python -m torch.distributed.run --master_port29501 --nproc_per_node4 train.py ...每次训练前也建议清理一下残留的训练进程pkill -f train.py5.2 batch size、学习率同步调整的经验值多卡训练不是把batch简单乘个卡数就行学习率也要跟着调整。YOLOv5本身有自动学习率调度会根据batch size推算合理的初始学习率但理解背后的规则仍然很重要方便手动控制和排查问题。线性缩放规则的思路是当batch size变为原来的k倍时学习率大致也乘以k。以默认的v5s配置为例官方在8卡batch 64环境下使用lr约0.01。如果只单卡跑batch 16学习率大致在0.002~0.003这个量级否则loss容易出现震荡或nan。实际中我一般先观察训练曲线loss发散就减半学习率loss下降缓慢就适当调大。更关键的是看warm-up阶段YOLOv5默认前3个epoch会从很小的学习率逐步升到设定值这个阶段不要因为loss下降太慢就中断训练给模型一点启动时间。多卡还有一个选项--sync-bn用于同步所有卡上的BatchNorm统计量。它的计算开销不小训练速度会下降。如果你的数据分布不一致、或检测头使用BatchNorm结构较多可以开启大多数场景下不需要尤其是显卡数量不多时收益小于损失。6. 容易被忽略的隐性开销日志、验证与软件习惯6.1 验证频率、日志写入与wandb的副作用很多人没有意识到训练过程中的验证val和日志同步也在消耗时间。如果每1个epoch都做全量验证验证集又比较大这部分耗时可能占整体时间的20%。YOLOv5默认在最后一个epoch做验证。如果你想在中间阶段看模型表现可以降低验证频率。较新版本支持--val参数指定验证间隔python train.py --data data.yaml --weights yolov5s.pt --img 640 --batch 16 --epochs 100 --val 5上面的命令表示每5个epoch验证一次既能看到中间结果又不会频繁打断训练。日志同步的坑更隐蔽。使用wandb或Comet等在线记录工具时网络不稳定会导致同步请求重试训练进程被卡住。我遇到过训练一小时有一半时间在等待日志上传的情况。如果只是自己调试完全没必要强上在线日志本地TensorBoard足够用。python train.py ... --project runs --name exp_test tensorboard --logdir runs训练过程中的checkpoint保存也会占用时间尤其模型较大时。快速验证阶段可以加--nosave跳过保存跑正式训练时再恢复。6.2 实用的监控与排查小工具最后分享几个我用得比较顺手的监控小工具能帮你在调优时快速判断问题出在哪里。持续输出GPU状态nvidia-smi dmon -f 60输出GPU利用率到CSV文件方便训练结束后复盘nvidia-smi --query-gpuutilization.gpu,memory.used,power.draw --formatcsv -l 60 gpu_log.csv查看CPU核心负载分布mpstat -P ALL 2如果数据增强是瓶颈CPU负载会异常高。针对CPU算力弱的机器YOLOv5的Mosaic增强是最大的开销来源。部分版本支持关闭或降低Mosaic概率在快速验证时可以临时禁用把时间留给真正的模型迭代。我个人的习惯是先跑一个最小实验确认代码和数据都没问题再开全量训练训练过程中偶尔用nvidia-smi撇一眼利用率如果利用率在95%以上就不折腾了如果掉得厉害就去查数据加载和验证频率。训练速度优化这件事最怕的就是凭感觉乱调参数。先量化再定位然后一次性改一个变量跑小数据集验证效果。这样几次下来你的训练流程会建立在自己的数据、自己的硬件基础上的最优解而不是照搬别人的默认配置。按照这个节奏单品缩短到原来的三分之一是完全可以做到的。
阅读完成 · 觉得有帮助？
