资讯详情

资讯详情

YOLOv5+PyCharm目标检测实战:从零搭建训练流水线

1. 从零搭建自己的目标检测训练流水线为什么选YOLOv5加PyCharm这套组合很多人第一次接触目标检测脑子里冒出来的第一个念头往往是“我是不是得先把深度学习、卷积神经网络、反向传播这些理论啃完才能动手”。我当初也是这么想的结果啃了半个月数学推导回头一看连一张图片都没跑通挫败感极强。后来我换了个思路——先跑通一个能出结果的完整流程再回头补理论。这个思路的转折点就是YOLOv5加PyCharm这套组合。YOLOv5这个目标检测框架最大的特点是把“训练自己的数据集”这件事的门槛压得非常低。它不像某些学术框架那样需要你手写数据加载器、自己定义损失函数、手动管理学习率调度它把这一整套流程都封装成了配置文件加命令行参数的形式。你只需要准备好图片、标注文件和一个描述数据集结构的yaml文件就能启动训练。而PyCharm作为Python的集成开发环境提供了代码补全、调试、虚拟环境管理、远程解释器等功能让你在调试数据预处理逻辑、查看训练脚本参数、排查报错时效率高很多。这套组合解决的核心问题是让一个有一定Python基础但不一定精通深度学习的人能够在自己的电脑或者租用的算力上用自己收集的图片数据训练出一个能识别特定目标的检测模型。适合的人群包括做工业质检的工程师想识别产品缺陷、做农业项目的同学想数果实、做安防相关Demo的开发者想检测特定行为、以及大量在校学生做课程设计或毕业设计需要落地一个视觉项目。我见过太多人卡在几个地方数据集格式不对导致训练报错、环境依赖版本冲突、训练到一半显存爆了、训练完了不知道怎么推理。这篇文章我会把这些环节全部拆开把每一步为什么这么做、容易在哪里翻车、怎么验证自己做对了都讲清楚。你跟着走一遍应该能少走至少两三个星期的弯路。2. 动手之前先把数据这件事想明白数据集采集与标注的取舍2.1 图片从哪来、要多少张才够用自己训练模型第一道坎不是代码是数据。我见过有人上来就问“我要训练一个识别安全帽的模型需要多少张图”这个问题没有标准答案但有一些经验区间可以参考。对于单一类别的目标检测如果目标特征比较明显比如安全帽颜色鲜艳、形状固定每个类别有300到500张标注图片通常就能得到一个初步可用的模型。如果目标形态多变、遮挡严重、背景复杂那可能需要1000到2000张甚至更多。我的建议是先标200张跑一轮训练看验证集上的表现再决定要不要继续加数据。不要一上来就标5000张万一标注规则有问题返工成本极高。图片的来源可以多样化自己用手机或相机拍摄、从公开数据集中筛选、用数据增强手段扩充。但要注意自己拍图时尽量覆盖不同的光照条件、拍摄角度、目标大小和背景环境。比如你要检测某个零件那就把这个零件放在不同颜色的桌面上拍、在不同时间段的光线下拍、从正面侧面俯视仰视都拍一些。这样训练出来的模型泛化能力才强。注意不要用同一段视频连续截帧来充当大量图片。这样得到的图片高度相似模型会“记住”而不是“学会”验证集指标虚高实际用的时候一塌糊涂。2.2 标注工具的选择与标注规范标注工具我用过几款比较顺手的是LabelImg和Labelme。LabelImg适合矩形框标注输出YOLO格式或Pascal VOC格式都支持操作简单快捷键熟练之后标注速度很快。Labelme更适合多边形标注如果你要检测的目标是不规则形状可以考虑它但输出格式需要额外转换。标注的时候有几个坑我踩过这里直接说结论框要贴紧目标边缘但不要切掉目标的一部分。框太大引入背景噪声框太小丢失目标信息。同一类别在不同图片中的标注标准要一致。比如“人”这个类别有的图你标了全身有的图只标了上半身模型会困惑。遮挡目标怎么标要提前定好规则。如果目标被遮挡超过70%我一般选择不标遮挡30%到70%之间按可见部分标框遮挡少于30%正常标全框。不要漏标。一张图里如果有5个目标你只标了3个模型会把没标的那2个当成背景来学习直接拉低召回率。标注完成后YOLOv5需要的标签格式是每张图片对应一个txt文件每行格式为类别编号 中心x坐标 中心y坐标 宽度 高度所有坐标都是相对于图片宽高的归一化值0到1之间。这个格式转换LabelImg可以直接导出不用自己写脚本。2.3 数据集划分与yaml配置文件标注完的图片和标签文件需要按比例分成训练集、验证集和测试集。常见比例是7:2:1或8:1:1。训练集用于更新模型权重验证集用于训练过程中监控模型表现和调整超参数测试集只在最后评估时用一次用来估计模型在真实场景下的表现。YOLOv5通过一个yaml文件来描述数据集。这个文件里要写清楚训练集和验证集的图片目录路径、类别数量、类别名称列表。我建议目录结构这样组织dataset/ images/ train/ val/ labels/ train/ val/对应的yaml文件内容大致如下train: dataset/images/train val: dataset/images/val nc: 3 names: [cat, dog, person]这里nc是类别数量names是类别名称列表顺序要和标注时的类别编号对应。这个文件放在YOLOv5项目根目录下的data文件夹里训练时通过参数指定。提示路径最好用相对路径这样整个项目文件夹拷贝到另一台机器上也能直接跑。如果用绝对路径换机器就要改配置容易忘。3. PyCharm里的工程化配置让训练脚本跑得稳、调得动3.1 虚拟环境与依赖安装的正确姿势PyCharm最实用的功能之一就是虚拟环境管理。我强烈建议为每个YOLOv5项目单独建一个虚拟环境不要用系统Python或者base环境。原因很简单YOLOv5对PyTorch、torchvision、numpy等库的版本有要求不同项目依赖不同版本时混在一起必然冲突。在PyCharm里创建虚拟环境的步骤打开项目后进入File Settings Project Python Interpreter点击齿轮图标选择Add然后选择Virtualenv Environment指定一个空文件夹作为环境目录基础解释器选你系统里的Python 3.8或3.9YOLOv5对这两个版本兼容性最好。创建完成后PyCharm会自动激活这个环境。接下来安装依赖。YOLOv5的根目录下有一个requirements.txt文件里面列出了所有需要的库。在PyCharm的Terminal里注意确认Terminal前面显示的是你刚创建的虚拟环境名称执行pip install -r requirements.txt这里有个坑requirements.txt里写的PyTorch版本可能不带CUDA支持如果你有NVIDIA显卡并且想用GPU训练需要单独安装带CUDA的PyTorch。去PyTorch官网找到对应你CUDA版本的安装命令比如pip install torch torchvision --extra-index-url https://download.pytorch.org/whl/cu118安装完成后在PyCharm的Python Console里执行import torch; print(torch.cuda.is_available())如果输出True说明GPU可用。3.2 用PyCharm的Run Configuration管理训练参数YOLOv5的训练入口是train.py通过命令行参数控制各种配置。在PyCharm里你可以为train.py创建一个Run Configuration把参数写在Parameters栏里这样每次训练不用手动敲命令点运行按钮就行。一个典型的训练参数配置如下--img 640 --batch 16 --epochs 100 --data data/mydata.yaml --weights yolov5s.pt --project runs/train --name exp1逐个解释这些参数的含义--img 640输入图片缩放到640x640。YOLOv5支持320到1280之间的值越大精度通常越高但速度越慢、显存占用越大。--batch 16每批处理16张图片。这个值受显存限制8G显存跑yolov5s模型大概能开到16跑yolov5m可能只能开到8。--epochs 100训练100轮。每轮会把训练集完整过一遍。--data data/mydata.yaml指定数据集配置文件路径。--weights yolov5s.pt从预训练权重开始训练这叫迁移学习能大幅加快收敛速度。yolov5s是最小的模型速度快yolov5m、yolov5l、yolov5x依次变大变慢但精度更高。--project runs/train训练结果保存的根目录。--name exp1本次实验的名称结果会保存在runs/train/exp1下。在PyCharm里配置好之后点击运行Terminal里会实时打印训练日志包括当前epoch、损失值、验证集mAP等指标。如果报错PyCharm的调试器可以让你在代码里打断点逐步排查问题这是比纯命令行训练方便的地方。3.3 训练过程中的显存与日志监控训练启动后最常遇到的问题就是显存不够报CUDA out of memory。解决办法按优先级排列减小--batch的值比如从16降到8或4。减小--img的值比如从640降到416或320。换更小的模型比如从yolov5m换成yolov5s。如果以上都不行检查是不是有其他程序占着显存用nvidia-smi命令查看。训练日志里要重点关注的指标指标含义正常表现box_loss边界框回归损失随训练逐渐下降obj_loss目标置信度损失随训练逐渐下降cls_loss分类损失随训练逐渐下降mAP0.5IoU阈值为0.5时的平均精度随训练逐渐上升mAP0.5:0.95多个IoU阈值下的平均精度随训练逐渐上升如果损失不降反升或者mAP一直上不去可能是学习率太大、数据标注有问题、或者模型容量不够。这时候需要回头检查数据和参数。4. 训练完了不算完推理验证与模型导出4.1 用验证集和测试集评估模型真实水平训练结束后runs/train/exp1目录下会生成一堆文件其中weights文件夹里有best.pt和last.pt两个权重文件。best.pt是验证集上表现最好的那一轮保存的权重通常用它来做推理。评估模型不能只看训练日志里的mAP还要用测试集跑一遍。YOLOv5提供了val.py脚本python val.py --weights runs/train/exp1/weights/best.pt --data data/mydata.yaml --img 640 --task test这会输出模型在测试集上的精确率、召回率、mAP等指标。如果测试集mAP明显低于验证集mAP说明模型过拟合了需要增加数据量、加数据增强、或者减小模型复杂度。4.2 单张图片和批量图片的推理推理用detect.py脚本python detect.py --weights runs/train/exp1/weights/best.pt --source test.jpg --img 640 --conf 0.25--source可以指定单张图片、整个文件夹、视频文件甚至摄像头。--conf 0.25是置信度阈值只有置信度高于0.25的检测框才会被保留。这个值调高会减少误检但可能漏检调低则相反。实际使用时根据场景需求调整。推理结果默认保存在runs/detect/exp目录下图片上会画出检测框和类别标签。我建议先用一批有代表性的测试图片跑一遍肉眼看看检测效果特别关注漏检和误检的情况。如果发现某类目标经常漏检可能是这类目标的训练样本太少如果背景被误检成目标可能是负样本不够。4.3 模型导出为部署格式训练好的.pt权重文件是PyTorch格式部署到不同平台时需要转换成其他格式。YOLOv5提供了export.py脚本python export.py --weights runs/train/exp1/weights/best.pt --include onnx支持的格式包括ONNX、TensorRT、CoreML、TFLite等。ONNX通用性最好可以在多种推理引擎上运行TensorRT在NVIDIA显卡上速度最快TFLite适合移动端和嵌入式设备。导出时要注意输入尺寸和动态轴设置。如果部署环境的输入尺寸固定导出时指定--img 640这样导出的模型只接受640x640的输入推理速度更快。如果需要动态尺寸加上--dynamic参数。注意导出ONNX后建议用onnxruntime加载模型跑一张测试图片确认输出结果和PyTorch版本一致。我遇到过导出后精度下降的情况排查发现是某些算子在不同框架下的实现有细微差异这时候需要调整导出参数或更换算子。5. 那些只有踩过才知道的坑常见报错与排查思路5.1 数据标注相关的报错报错AssertionError: train: No labels found这个报错的意思是训练集目录下没有找到标签文件。排查步骤确认data/mydata.yaml里的train路径指向的是包含图片的文件夹YOLOv5会自动在同级目录下找同名的labels文件夹。比如图片在dataset/images/train标签应该在dataset/labels/train。如果目录结构不对要么调整目录要么在yaml里分别指定图片和标签路径。报错IndexError: list index out of range在加载标签时通常是标签文件里有空行或者格式不对。YOLOv5要求每行5个值用空格分隔。检查标签文件是否有空行、是否有非数字字符、坐标值是否在0到1之间。写个小脚本批量检查一遍所有标签文件把不合规的找出来修正。5.2 环境与版本相关的报错报错ImportError: cannot import name xxx from torch这是PyTorch版本和YOLOv5代码不匹配。YOLOv5的不同版本对PyTorch版本有要求去YOLOv5的GitHub页面看对应版本的requirements.txt严格按照里面写的版本安装。不要盲目升级到最新版。报错RuntimeError: CUDA error: no kernel image is available for execution on the device这是PyTorch的CUDA版本和显卡驱动不匹配。用nvidia-smi查看驱动支持的CUDA版本然后安装对应版本的PyTorch。比如驱动支持CUDA 11.8就装cu118版本的PyTorch。5.3 训练过程中的异常训练损失正常下降但mAP一直是0这种情况通常是类别编号和yaml里的names顺序对不上。比如标注时把“猫”标为0、“狗”标为1但yaml里写的是names: [dog, cat]那模型学到的类别就全反了。检查标注时的类别映射关系确保和yaml一致。训练到一半突然中断没有报错信息大概率是显存溢出被系统杀掉了或者电脑休眠了。检查系统日志确认是不是OOM。如果是减小batch size。另外训练时把电脑的休眠关掉或者用nohup命令在后台跑。验证集mAP波动很大时高时低可能是验证集太小指标不稳定。建议验证集至少要有100张图片。另外检查数据增强是否过于激进导致验证集图片也被过度变换。YOLOv5默认只在训练时做数据增强验证时不做但如果你改了代码就可能影响。6. 从能跑到好用提升模型效果的几个实操方向6.1 数据增强策略的调整YOLOv5默认开启了Mosaic、HSV增强、随机翻转等数据增强。Mosaic把4张图片拼成1张能显著提升小目标检测效果但如果你的目标都很大Mosaic可能反而有害。可以在--hyp参数指定的超参数文件里调整增强强度或者直接关掉某些增强。我的一般做法是先用默认增强跑一轮看效果如果过拟合严重加大增强如果欠拟合减小增强。超参数文件在data/hyps目录下复制一份改参数训练时用--hyp指定。6.2 锚框的重新聚类YOLOv5默认的锚框是基于COCO数据集聚类的如果你的目标尺寸分布和COCO差异很大重新聚类锚框能提升检测精度。YOLOv5提供了utils/autoanchor.py脚本训练时加上--noautoanchor关闭自动锚框然后手动运行聚类脚本把得到的锚框值写进模型配置文件。6.3 学习率与优化器的选择默认使用SGD优化器学习率通过余弦退火调度。如果训练不稳定可以换Adam优化器加上--optimizer Adam参数。学习率初始值默认是0.01如果损失震荡厉害降到0.001试试。这些参数都在超参数文件里改完重新训练即可。6.4 模型剪枝与量化如果最终要部署到算力有限的设备上可以对训练好的模型做剪枝和量化。剪枝去掉不重要的神经元量化把浮点权重转成整数都能减小模型体积、加快推理速度。YOLOv5官方没有直接提供剪枝脚本但可以用PyTorch的torch.quantization做量化或者用第三方工具如NNCF做剪枝量化。这部分操作复杂度较高建议先确保模型精度达标再考虑。7. 一些关于训练节奏和心态的体会我刚开始做自己数据集训练的时候总想一次就把所有参数调到最优结果在电脑前坐了一整天改了十几版配置最后连一个能用的模型都没训出来。后来我学乖了先用小数据量、小模型、少轮次快速跑通全流程确认数据格式、环境配置、训练脚本都没问题再逐步加数据、换大模型、调参数。这个思路帮我省了大量时间。另外不要迷信网上的“最佳实践”。别人的最优参数是建立在别人的数据集和硬件环境上的直接搬过来大概率水土不服。你要做的是理解每个参数的作用然后根据自己的情况做小步调整。每次只改一个参数观察指标变化记录下来。这样积累的经验才是你自己的。最后说一个容易被忽略的点保存训练日志和配置文件。YOLOv5每次训练会在runs/train/exp目录下保存opt.yaml里面记录了本次训练的所有参数。如果你训了很多轮回头想复现某个效果好的版本没有这个文件就只能靠回忆。我现在的习惯是每次训练完把opt.yaml和results.csv单独拷贝到一个归档文件夹按日期和实验名命名方便以后查阅。训练模型这件事三分靠算法七分靠数据。把数据质量提上去把标注规范定清楚比调任何参数都管用。
觉得有用,分享给同行:

为您的企业打造数字门面

稳重轻奢商务风格,端正雅致视觉,长效耐看不易过时。

立即咨询 →