ARTICLE DETAIL

资讯详情

深耕网站建设、视觉设计与SEO优化的一线实战洞察。

YOLOv5目标检测实战:从环境配置到训练自己的数据集全流程

YOLOv5目标检测实战:从环境配置到训练自己的数据集全流程 1. 先聊聊从入门到入狱这个标题YOLOv5到底学的是什么1.1 一次目标检测任务让我决定彻底搞懂它我最初接触到 YOLOv5 是因为一个很实际的需求要给厂里一条流水线上的工件做数量统计和位置标定。当时团队里有人提议用传统图像处理但零件种类一多、光照一变阈值怎么调都崩。后来有人甩过来一句试试 YOLOv5 吧我就从一个只会 pip install 的小白硬着头皮走上了自学目标检测的路。这篇笔记就是我从零开始跑通 YOLOv5 完整流程的记录环境配置、官方权重推理、标注自己的数据集、训练、评估、导出模型。如果你也是刚接触深度学习、想在真实项目里用 YOLOv5 训练自己的数据集这篇笔记应该能帮你少走不少弯路。标题里从入门到入狱是句调侃意思是 YOLOv5 功能确实强大但乱用、滥用、错用轻则训练出一堆废模型浪费算力重则把不合格的模型放到不该放的场景里那真就离出事不远了。1.2 功能强大不要乱用不是开玩笑先说清楚 YOLOv5 能做什么。它是一个单阶段目标检测模型核心任务就是在图片里找到物体并框出来可以检测人、车、动物、缺陷、工件、农作物几乎任何你手头有图片数据的东西。它速度快、精度高、生态成熟网上教程一抓一大把看起来轻轻松松就能训练一个自己的检测器。但正因为门槛低才更容易出问题。有人拿它做人脸识别相关的敏感应用不考虑合规性有人拿它训练自动抢票、作弊外挂之类的工具这已经不只是技术滥用的问题了有人完全不懂数据集怎么来爬了一堆版权不明的图片就开训模型上线就被告还有人拿 5 张图就敢训生产级模型结果误检率奇高直接背锅。我入狱这两个字不是真的指坐牢而是想表达如果你只学会了操作、没学会判断什么场景能用、什么数据能训、什么后果要承担那翻车是迟早的事。所以这篇笔记不光是教你怎么跑通流程也会在关键节点提醒你哪些地方别乱来。1.3 这篇笔记覆盖什么、不覆盖什么覆盖的内容包括YOLOv5 环境搭建CPU 和 GPU 都讲官方权重推理图片、视频、摄像头实时检测用 labelImg 标注自己的图片数据组织数据集目录结构、编写 data.yaml 配置文件完整跑一次训练并理解核心训练参数评估模型效果、导出 ONNX 模型我在实操中遇到过的几个坑和排查思路不覆盖的内容包括YOLOv5 内部每一层的详细数学推导、自定义网络结构修改、分布式多机训练、模型量化部署到边缘设备的全流程。这些是进阶内容后面我的笔记系列里可能会单独写。2. 环境准备先把地基打牢后面才不返工2.1 硬件与软件版本匹配是最大的坑我第一次装 YOLOv5 环境照着网上的教程一顿操作结果 torch 版本和 CUDA 版本对不上import torch 时说 CUDA unavailable折腾了整整一个晚上。后来我总结出一条经验先确定你机器的硬件再倒推装什么版本的驱动、CUDA、PyTorch顺序反了就是无限踩坑。先看你的显卡NVIDIA 显卡建议用 GPU 训练速度比 CPU 快几十倍不止没有 NVIDIA 显卡或者用 Mac也能跑但只能用小模型、小数据集做学习验证训练 100 轮的 COCO128 都要等很久纯 CPU 环境强烈建议先装 CPU 版 PyTorch别去折腾 CUDA。一个我自己试过的稳定组合是这样的组件版本推荐备注Ubuntu20.04 / 22.04Windows 也能跑但坑更多Python3.8 3.10YOLOv5 官方测试过的主要版本PyTorch1.13 或 2.x具体看 CUDA 版本CUDA11.7 或 11.8不是越新越好要匹配 PyTorchcuDNN对应 CUDA 版本一般装 PyTorch 时自带显卡驱动对应 CUDA 要求的版本用 nvidia-smi 查看判断 CUDA 和 PyTorch 是否匹配的最快方法去 PyTorch 官网的 Get Started 页面选择你的系统、安装方式、CUDA 版本它会直接生成对应的安装命令。不要自己去网上乱找 pip 命令。2.2 使用 conda 创建独立环境我强烈建议用 conda 创建一个独立的虚拟环境不要直接装在系统 Python 里。原因是 YOLOv5 的依赖项跟其他项目经常冲突独立环境出问题了直接删掉重来不污染系统。# 创建环境 conda create -n yolov5 python3.9 -y # 激活环境 conda activate yolov5 # 安装 PyTorch以 CUDA 11.8 为例 pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 # 验证 PyTorch 是否能用 GPU python -c import torch; print(torch.__version__, torch.cuda.is_available())如果输出2.x.xcu118 True说明 GPU 版本可用。如果显示False先别急着往下走检查驱动版本和 CUDA 是否装对。2.3 拉取 YOLOv5 源码并安装依赖YOLOv5 的代码托管在 GitHub 上使用 git 拉取git clone https://github.com/ultralytics/yolov5 cd yolov5 pip install -r requirements.txt这里有个细节requirements.txt里安装的依赖包含torch和torchvision。如果你之前已经手动装了 GPU 版 PyTorch建议在安装 requirements 时加一个--no-deps参数跳过重装或者直接注释掉 requirements 里的 torch 那几行否则 pip 可能会把 GPU 版覆盖成 CPU 版。这是我踩过的第二个坑安装完发现 torch 变成 CPU 版了。验证环境是否正常python detect.py --source data/images/bus.jpg --weights yolov5s.pt如果能在runs/detect/exp下看到带框的 bus 图片说明整个环境已经通了。注意第一次运行会自动下载yolov5s.pt权重文件大概 14MB 左右如果下载慢可以考虑手动下载放到项目根目录。3. 第一天实战用官方权重跑通检测流程3.1 官方权重怎么选YOLOv5 官方提供了 n/s/m/l/x 五个尺度的预训练权重全部是在 COCO 数据集上训练好的能检测 80 类常见物体人、汽车、猫、狗、杯子、椅子等。模型参数量模型大小推理速度精度(mAP)适用场景YOLOv5n1.9M约 4MB最快较低嵌入式/移动端YOLOv5s7.2M约 14MB快中等一般实时检测YOLOv5m21.2M约 41MB中等较高精度优先YOLOv5l46.5M约 90MB较慢高服务器端YOLOv5x86.7M约 170MB慢最高追求极限精度第一次跑通流程建议用yolov5s.pt速度、精度、显存占用都适中。我自己的经验是不要一上来就追求 x 模型先用 s 把全流程跑通后面再换大模型提升精度这样排查问题会容易很多。3.2 detect.py 参数逐个拆解detect.py是 YOLOv5 的推理脚本核心参数如下python detect.py \ --weights yolov5s.pt \ # 权重文件路径 --source data/images/bus.jpg \ # 输入图片 --conf-thres 0.25 \ # 置信度阈值低于该值的目标会被过滤 --iou-thres 0.45 \ # NMS 的 IoU 阈值控制重叠框的合并 --img 640 \ # 推理尺寸 --device 0 \ # GPU 设备号CPU 用 --device cpu --save-txt \ # 同时保存 txt 标注文件 --classes 0 2 \ # 只检测指定类别COCO 的类别索引参数含义不难理解但有两个容易被忽略的点--conf-thres设太低比如 0.1会产生大量误检框设太高比如 0.8会漏检。实际项目中应该先在测试集上跑一遍观察置信度分布再定阈值。--img影响检测精度和速度的平衡。输入尺寸越大小目标越容易检到但推理越慢。一般 640 是默认值如果小目标很多可以试试 1280显存要相应加大。--source参数非常灵活支持图片、目录、视频、摄像头、网络流等各种输入。传 URL 时要用引号包住python detect.py --weights yolov5s.pt --source https://ultralytics.com/images/zidane.jpg3.3 从单张图片到视频流跑通单张图片后可以试试视频和摄像头# 视频文件 python detect.py --weights yolov5s.pt --source test.mp4 # 摄像头0 表示默认摄像头 python detect.py --weights yolov5s.pt --source 0视频检测的本质是逐帧推理CPU 上跑 640 尺寸的视频帧率会很低大概只有几帧每秒。GPU 上跑会好很多但也到不了实时。如果要做实时摄像头检测建议用yolov5n或yolov5s加上--half参数开启 FP16 推理。这里有个体验上的小建议第一次跑视频时输出结果会保存在runs/detect/expN目录里每张图片/每帧都不会覆盖之前的输出结果这是 YOLOv5 的一个细节设计——每次运行自动新建 exp 目录方便对比不同参数的效果。4. 从会用到懂用YOLOv5 结构和原理速览4.1 网络骨架、颈部、检测头的关系这个部分我不打算堆公式但完全不懂原理去调参跟摸黑走夜路没区别。YOLOv5 的网络结构可以理解成三个部分Backbone骨干网络、Neck颈部、Head检测头。BackboneCSPDarknet负责提取图像特征。图片输入后会经过一系列卷积和下采样从原始像素变成多尺度的特征图越往后特征图越小但语义信息越丰富NeckPANet负责把不同尺度的特征融合。小物体需要浅层的细节信息大物体需要深层的语义信息PANet 通过自顶向下和自底向上的路径把两者结合起来Head负责输出最终结果包括每个预测框的位置、尺寸、类别概率和置信度。YOLOv5 的输出有三个尺度对应检测大、中、小三种目标。如果你用export.py导出模型后再用 Netron 查看 ONNX 图能看到三个输出头这就是为什么 YOLOv5 对小目标有一定检测能力的原因。4.2 三种尺寸模型的选择逻辑很多新手会问训练自己的数据集到底选哪个规模我的建议是如果数据集比较小几百张到一两千张优先用 s 或 n。模型参数量越大过拟合的风险越高如果数据集充分5000 张以上且目标多是中小物体可以上 m 或 l先跑通一次完整训练用 s 模型看 loss 能不能收敛、验证集的指标大概什么水平再决定要不要换更大的模型。用 x 模型跑小数据集结果往往不如 s 模型这就是典型的功能强大不要乱用——不是所有场景都需要大模型。4.3 Mosaic 数据增强为什么在训练里至关重要训练 YOLOv5 时默认开启 Mosaic 增强原理是把 4 张训练图片随机缩放后拼成一张新图片等于用一张图同时让模型看到更多目标、不同尺度、不同背景。这能显著提升小目标检测能力也让模型对遮挡更鲁棒。但 Mosaic 也不是没缺点。如果数据集本身分布不均匀比如某类目标特别多拼图可能会让模型学到错误的类别关联。我在一个项目中就发现Mosaic 增强导致模型把座椅和人绑定了后来在hyp.scratch-low.yaml里把mosaic从 1.0 降到 0.5 才改善。4.4 训练时数据被喂了几遍是什么意思初学时会经常看到epoch这个词它指的是整个训练集被完整遍历的次数。比如你有 100 张训练图片batch size 设为 16那么一个 epoch 就是把这 100 张图片分 7 批100/16向上取整喂给模型。注意epoch和batch是两个经常被混淆的概念epoch完整过一遍数据集batch size一次喂给模型多少张图iterations/step一个 epoch 内更新的次数等于训练集图片数/batch size向上取整。YOLOv5 的日志里显示的 Epoch、GPU_mem、box_loss、obj_loss、cls_loss 等指标前 50 个 epoch 主要看 loss 是否在下降后面再看验证集的 mAP 是否上升。一次训练几百轮等于数据被喂了几百遍模型会逐渐记住数据里的特征模式。5. 训练自己的数据集绕不开的脏活累活5.1 数据采集与标注工具选择自己训练前先要搞清楚你的数据从哪里来这个环节最容易被忽略但数据质量直接决定模型上限。数据采集的原则尽量在真实部署场景中采集模拟场景训出来的模型到真实场景往往精度暴跌覆盖尽量多的情况不同光照、不同角度、不同距离、不同遮挡、不同背景每个类别至少 1000 张以上简单场景可以少一点复杂场景建议更多太少的话模型很容易过拟合类别不要一开始就搞很多建议先做 1 到 3 类的试点项目跑通后再扩展。标注工具推荐两个工具优点缺点labelImg经典、轻量、无需联网界面老式但够用X-AnyLabeling支持半自动标注用检测模型预标注依赖较多要装模型我平时主要用 labelImg。安装简单标注完保存为 YOLO 格式会自动生成与图片同名的.txt文件。pip install labelImg labelImg # 打开后记得把保存格式切换为 YOLO标注 BTW 几个操作技巧先设置默认保存目录自动保存否则标完 50 张忘了保存就白干了打开自动保存和单类模式可以大幅加快单类别目标的标注速度框尽量紧贴目标边缘但不要把关键特征切掉。5.2 数据集目录结构与 data.yaml 配置YOLOv5 训练要求的数据集目录结构如下dataset/ ├── images/ │ ├── train/ # 训练集图片 │ └── val/ # 验证集图片 ├── labels/ │ ├── train/ # 训练集标签txt │ └── val/ # 验证集标签txt └── data.yaml # 数据集配置文件注意图片和标注文件的文件名必须一一对应。比如0001.jpg对应0001.txt拼错一个都会导致训练时报no labels found。我一般按照 8:2 或者 9:1 的比例划分训练集和验证集。测试集可以另外留一份但 YOLOv5 官方的流程通常只用 train 和 val。data.yaml 的内容# 数据集根目录 path: /home/user/dataset train: images/train val: images/val # 类别数量 nc: 2 # 类别名称顺序与标注文件里的 class_id 一一对应 names: 0: cat 1: dog这里有一个很容易踩的坑class_id 必须从 0 开始连续编号。如果你在标注时只有 3 个类别但 class_id 是 1、2、3训练会直接报错。标注文件 txt 的每一行格式是class_id x_center y_center width height后面四个值都是归一化坐标即像素坐标除以图片宽高取值范围 0 到 1。labelImg 会自动帮你算好所以正常标注生成的文件不需要手改。5.3 train.py 参数详解与一次完整训练训练命令长这样python train.py \ --data data.yaml \ --weights yolov5s.pt \ --epochs 100 \ --batch-size 16 \ --img 640 \ --device 0 \ --name my_custom_model参数背后的逻辑--weights yolov5s.pt基于官方预训练权重做迁移学习。千万不要用--weights 从零开始训练除非你有海量数据和充足算力否则收敛极慢、精度还低--epochs 100训练多少轮。100 轮是一个比较合理的起点如果你的数据量很小50 轮可能就够了数据量大或者任务复杂可以到 200 轮、300 轮--batch-size 16显存不够就调小到 8 或 4显存够用可以调大到 32 甚至 64。batch-size 越大训练越稳定但不要超过显卡显存上限--img 640训练时图片缩放的目标尺寸。训练图片会先按比例缩放到短边 640再做 padding 补到 640×640--cache把图片缓存到内存中加快数据读取如果你的内存足够比如 32GB 以上建议加上。完整训练一次控制台日志里会有这些信息Epoch gpu_mem box obj cls labels img_size 1/100 6.21G 0.0893 0.0205 0.0013 52 640 2/100 6.21G 0.0714 0.0182 0.0011 48 640重点关注box框损失、obj目标损失、cls分类损失的变化趋势。正常情况下这些值应该随着 epoch 增加而波动下降。如果 loss 不降反升或者震荡特别剧烈说明学习率太高、batch-size 太小或者数据有问题。训练结束后结果存放在runs/train/name/目录下包括weights/best.pt验证集上 mAP 最高的权重部署时首选weights/last.pt最后一轮的权重一般用于续训results.pngloss、mAP、PR 曲线等训练过程图像confusion_matrix.png混淆矩阵查看类别的互相混淆情况train_batch*.jpg训练时的 batch 数据增强预览图能直观看到数据增强效果。5.4 超参数调整的逻辑先别急着动YOLOv5 的超参数配置在data/hyps/hyp.scratch-low.yaml里包含lr0初始学习率、lrf最终学习率因子、momentum、weight_decay、warmup_epochs、mosaic、hsv_h等一堆参数。新手最容易犯的错是一开始就乱调超参数。我的建议是第一次训练全部用默认参数目的是验证数据、代码流程没有 bug如果 loss 收敛但 mAP 不够高优先从数据层面改进增加数据量、标注质量、类别均衡而不是调超参如果模型过拟合训练 loss 很低但验证 mAP 上不去再考虑调整weight_decay、dropout、mosaic、mixup等正则化相关参数如果训练震荡不稳再把lr0调小一半试试。我曾经遇到过一次极端情况某个类别只有 20 张图另一个类别有 2000 张图训练出来那个少样本类别 mAP 只有 0.05。后来加了数据增强和类别权重才勉强到 0.4。这让我深刻体会到超参数是锦上添花数据分布才是决定模型上限的根本。6. 训练完成之后验证、导出与部署6.1 test.py 或 val.py 的指标怎么看训练结束后YOLOv5 会自动在验证集上评估一次日志里会显示Class Images Instances P R mAP50 mAP50-95 all 200 300 0.923 0.887 0.934 0.712这几个指标的含义PPrecision精确率模型检测出的框中真正是目标的比例。P 高说明误检少RRecall召回率真实目标中被模型检测出来的比例。R 高说明漏检少mAP50IoU 阈值为 0.5 时的平均精度最常用的粗粒度指标mAP50-95IoU 阈值从 0.5 到 0.95 的平均精度要求框的定位更精确。实际项目里P 和 R 要根据业务场景权衡。比如质检场景宁可误检也不能漏检就把置信度阈值调低而安防告警场景希望减少误报就把阈值调高。如果对验证集上的表现不满意可以回到训练阶段但不要指望靠调参数起死回生。先去看confusion_matrix.png查清楚是类别混淆、漏检还是定位不准再有针对性地改数据。6.2 导出 ONNX 并部署到自己的程序里训练好的best.pt是 PyTorch 格式实际部署时往往需要转成 ONNX 或者 TensorRT 等格式。导出 ONNX 的命令python export.py --weights runs/train/my_custom_model/weights/best.pt --include onnx导出的.onnx文件可以用 ONNX Runtime 加载在 Python 里部署import onnxruntime as ort import cv2 import numpy as np sess ort.InferenceSession(best.onnx) input_name sess.get_inputs()[0].name output_names [o.name for o in sess.get_outputs()] img cv2.imread(test.jpg) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img cv2.resize(img, (640, 640)) img img.astype(np.float32) / 255.0 img np.transpose(img, (2, 0, 1))[None] outputs sess.run(output_names, {input_name: img})导出后要做一个重要验证ONNX 模型的推理结果要和 PyTorch 模型保持一致。因为导出过程中可能有算子兼容问题导致结果对不上。用同一张测试图分别跑.pt和.onnx对比检测框坐标和置信度误差在合理范围内才算部署成功。如果要在 Jetson Nano 或者手机这类边缘设备上跑可以进一步导出 TensorRT 的 engine 文件速度比 ONNX 快不少。这是另外一个话题后面有机会单独写。7. 我在实操中踩过的坑附完整排查思路7.1 CUDA out of memory不是无脑调小 batch训练时报错RuntimeError: CUDA out of memory这是新手最常遇到的错误。很多人第一反应是把 batch-size 减半然后接着跑但这只是治标。我的排查顺序是这样的先用nvidia-smi看显存占用情况。如果是其他进程占用了显存先 kill 或者换卡这个最常见如果自己的训练就爆显存了优先把--img从 640 降到 512 或 416这个比调 batch-size 更有效因为显存占用和输入尺寸是平方关系batch-size 从 16 降到 8、4添加--cache-images可能会导致内存不足如果内存吃紧就不要加实在不行可以考虑梯度累积但这需要改代码逻辑不太适合新手。7.2 训练 loss 不下降问题不在网络在数据有次我训练一个识别果品的模型50 个 epoch 过去了box_loss 一直在 0.1 附近横盘val mAP 永远都是 0。排查过程是这样的先看训练日志发现 labels 数量为 0训练时一个目标框都没读进去。问题出在我把标签文件夹路径写错了data.yaml 里 val 配到了空目录修好路径后重新训练loss 还是降不下去。用runs/train/exp/train_batch0.jpg查看训练图像发现标注框和目标位置完全对不上——原来是 labelImg 保存时图片路径变了标注文件和图片不是同一张。最终我发现训练时--data指向的路径用的是相对路径而我在不同目录下执行命令导致图片找不到部分训练图片实际上没有对应的有效标签。后来我把所有的路径都改成绝对路径并把一张张图片和标签重新对照了一遍损失曲线才恢复正常。所以如果你的 loss 不降第一件事永远都是打开train_batch0.jpg看标注对不对、是不是漏读标签了。别怀疑模型结构90% 的概率是你数据流程有问题。7.3 精度高但检测错位小目标丢失与Anchor聚类另一个头疼的问题是mAP 看起来不错但实际用起来小目标经常丢。举个例子检测远处的人在 640 输入尺寸下一张 1920×1080 的图里人可能只有 30×60 像素模型很难抓住。我的处理方案是把推理尺寸从 640 提高到 1280效果立竿见影但显存和速度代价不小使用切片推理--slide-overlap把大图切成小块重叠推理再拼接结果检查数据集中小目标的比例。如果小目标太少模型学不到对应尺度的特征可以考虑做过采样复制或者专门用小目标样本微调一个 epoch。YOLOv5 默认的 Anchor 是基于 COCO 数据集计算的不一定适合你的目标尺寸分布。训练前可以用工具计算自己数据集的 Anchor 并配置到模型里但这一步其实在 YOLOv5 里默认开启自动锚框计算一般不用手动干预。只有在目标尺寸分布和常规相差特别大时才需要关注。8. 写在最后能力越大责任越大回到标题那个入狱梗。YOLOv5 本身是一把工具就像一把刀可以切菜也可以伤人。学习它的过程里我觉得最重要的不是学会多少参数、调通多少模型而是建立起对技术的判断力和边界感。具体来说我给自己定了三条规矩数据必须有明确来源和授权。不爬不明来源的图片不把他人版权数据偷偷拿来训练商业模型应用场景要负责任。涉及人脸、隐私、公共安全等敏感场景先跟相关人员确认合规性不要一个人闷头开发完就上线明确知道模型的边界。模型不是万能的误检漏检在所难免。在需求沟通和方案设计时把模型的适用条件和失败可能性说清楚比什么都重要。这篇笔记是自学笔记系列的第一篇挺长的能看到这里说明你真的想学会看不下去也很正常实操一次比读十遍都管用。下一步你可以拿一个 100 张图片的小数据集跟着完整流程走一遍遇到问题再回来对照笔记排查。我个人在实际操作中最想再强调一次的是第一次训练老老实实用默认参数目标是跑通流程不是追求精度。等全流程走通了你再回去看那些 mAP 指标、超参数调优的文章会有完全不同的体会。YOLOv5 入门不难难的是把每个环节做得扎实。数据、标注、训练、验证、部署每一步都值得你反复打磨。
返回列表