ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

Yolov5+Pytorch训练自己数据集:从环境搭建到模型部署全流程

Yolov5+Pytorch训练自己数据集:从环境搭建到模型部署全流程 简介这份资源面向希望上手目标检测的深度学习学习者与工程实践者围绕Yolov5与PyTorch框架完整演示如何训练自己的自定义数据集。内容覆盖数据标注与预处理、配置文件修改、训练脚本运行、指标监控以及测试部署等关键环节并延伸至数据增强、模型微调与多GPU训练等进阶话题适合具备一定Python基础、想从零跑通检测项目的读者。压缩包共69个文件约13.81MB以20个Python脚本和9个YAML配置为主辅以JPG、PNG等示例图片、Markdown说明文档、Shell脚本、Notebook及Dockerfile兼顾代码、配置与容器化环境搭建。目前已有302人学习下载。通过这份实战资料读者可对照完整目录结构理解Yolov5项目组织方式掌握从数据准备到模型评估的落地流程并借助示例图片与权重脚本快速验证训练效果减少自行摸索的排错成本。1. 从一堆标注图片到能跑的检测模型Yolov5Pytorch 训练自己数据集到底难在哪手里有几百上千张标注好的图片想训练一个能识别自己业务目标的目标检测模型这件事听起来门槛不高但真正动手的人大多会在环境配置、数据格式、超参数这三道坎上卡住。Yolov5 配合 Pytorch 是目前工业界落地最成熟的组合之一源码结构清晰、训练速度快、部署链路完整从 yolov5 源码拉下来到跑通第一个 epoch熟练的话半小时内能搞定但第一次接触的人往往要折腾一整天。这篇内容面向的是手上有自己的数据集、想用 Yolov5 训练出可用检测模型的工程师和算法爱好者不聊空洞的算法推导只讲从环境搭建、数据组织、配置修改到训练调参、推理验证的完整可复现路径。目标检测这个方向近几年迭代很快yolov8、yolo26 等新版本不断出现但 Yolov5 依然是很多团队的首选基线原因很简单文档全、踩坑记录多、部署工具链成熟遇到问题能搜到答案。下面按实际操作的顺序把每一步的命令、参数和判断标准讲清楚。2. 环境搭建与 Yolov5 源码准备把地基打牢再动工2.1 Pytorch 环境配置的版本对应关系环境配置是第一个翻车高发区。Pytorch 和 Python、CUDA 驱动之间有严格的版本对应关系装错了轻则跑不了 GPU重则 import 直接报错。我一般推荐用 conda 建独立环境避免和系统 Python 打架。先确认显卡驱动支持的 CUDA 版本nvidia-smi输出右上角的CUDA Version是驱动支持的最高 CUDA 版本比如显示 12.1那你装的 Pytorch 对应的 CUDA 版本不能超过 12.1。然后建环境conda create -n yolov5 python3.9 -y conda activate yolov5Python 版本选 3.8 到 3.10 之间最稳3.11 以上部分依赖包可能没有预编译 wheel。接下来装 Pytorch去 Pytorch 官网的安装命令生成器选对应版本比如 CUDA 11.8 的组合pip install torch2.0.1 torchvision0.15.2 --index-url https://download.pytorch.org/whl/cu118装完验证import torch print(torch.__version__) print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0))cuda.is_available()返回 True 才算 GPU 可用。如果返回 False先检查是不是装成了 CPU 版本再看驱动版本是否匹配。这一步没有后悔药装错了就老老实实卸载重装别想着凑合。2.2 拉取 Yolov5 源码与依赖安装源码直接从官方仓库克隆注意选稳定分支而不是 mastergit clone -b v7.0 https://github.com/ultralytics/yolov5.git cd yolov5 pip install -r requirements.txtrequirements.txt 里包含了 opencv、matplotlib、tqdm、pyyaml 等依赖。安装过程中如果 opencv 编译报错通常是缺少系统库Ubuntu 下补一句sudo apt-get install -y libgl1-mesa-glx libglib2.0-0装完后跑一个自检python detect.py --source data/images/bus.jpg --weights yolov5s.pt这条命令会自动下载 yolov5s 的预训练权重然后对示例图片做推理输出保存在runs/detect/exp/下。能正常输出带框的图片说明环境和源码都没问题。这一步很关键很多人跳过自检直接上自己的数据结果出了问题分不清是环境问题还是数据问题。2.3 目录结构与关键文件说明Yolov5 的目录结构不复杂但几个关键文件必须心里有数路径作用data/存放数据集配置 yaml 和示例图片models/网络结构定义 yaml 和公共模块utils/数据加载、损失计算、指标评估等工具train.py训练入口脚本detect.py推理脚本runs/训练和推理的输出目录自己训练时主要改两个地方data/下新建数据集配置文件models/下选一个模型结构 yaml。其他文件基本不动。3. 数据集组织与标注格式转换让 Yolov5 读懂你的数据3.1 目标检测标注工具与 YOLO 格式要求目标检测常用的标注工具就那几个LabelImg、Labelme、CVAT、Roboflow。LabelImg 最轻量适合小数据集快速标CVAT 适合团队协作和视频标注。不管用哪个工具最终要转成 YOLO 格式。YOLO 格式的标注是一个 txt 文件对应一张图片每行一个目标格式为class_id x_center y_center width height其中坐标都是归一化到 0 到 1 之间的浮点数相对于图片宽高。比如一张 640x480 的图里有个框在 (100, 80) 到 (300, 240)转换后0 0.3125 0.3333 0.3125 0.3333计算方式x_center (100300)/2/640 0.3125y_center (80240)/2/480 0.3333width (300-100)/640 0.3125height (240-80)/480 0.3333。如果标注工具输出的是 VOC 格式的 xml需要转换。写个脚本批量处理import xml.etree.ElementTree as ET import os from PIL import Image def voc_to_yolo(xml_path, img_path, class_map, output_dir): tree ET.parse(xml_path) root tree.getroot() img Image.open(img_path) w, h img.size lines [] for obj in root.iter(object): cls_name obj.find(name).text if cls_name not in class_map: continue cls_id class_map[cls_name] bbox obj.find(bndbox) x1 float(bbox.find(xmin).text) y1 float(bbox.find(ymin).text) x2 float(bbox.find(xmax).text) y2 float(bbox.find(ymax).text) xc (x1 x2) / 2 / w yc (y1 y2) / 2 / h bw (x2 - x1) / w bh (y2 - y1) / h lines.append(f{cls_id} {xc:.6f} {yc:.6f} {bw:.6f} {bh:.6f}) txt_name os.path.splitext(os.path.basename(xml_path))[0] .txt with open(os.path.join(output_dir, txt_name), w) as f: f.write(\n.join(lines))class_map是类别名到 id 的映射字典必须从 0 开始连续编号。转换后要抽查几张用可视化脚本画框确认坐标没算错。我见过有人转换时忘了归一化训练 loss 直接飞到天上。3.2 数据集划分与 yaml 配置文件编写数据集按 8:1:1 或 7:2:1 划分训练集、验证集、测试集。目录结构建议这样组织dataset/ images/ train/ val/ test/ labels/ train/ val/ test/然后在data/下新建mydata.yamlpath: /home/user/dataset train: images/train val: images/val test: images/test nc: 3 names: [person, car, dog]nc是类别数names按 id 顺序列出类别名。注意path写绝对路径最稳相对路径容易因为工作目录不同而出错。train、val、test是相对于path的子路径。3.3 数据增强策略与缓存机制Yolov5 默认开启了 mosaic、mixup、HSV 增强、随机翻转等策略。mosaic 把四张图拼成一张对小目标检测效果提升明显但如果你数据里目标都很大mosaic 可能反而引入噪声。可以在data/hyps/hyp.scratch-low.yaml里调整mosaic: 1.0 # 关闭设为 0 mixup: 0.0 # 默认关闭 hsv_h: 0.015 hsv_s: 0.7 hsv_v: 0.4 flipud: 0.0 fliplr: 0.5首次训练时 Yolov5 会生成标签缓存文件train.cache和val.cache放在 labels 目录旁边。如果改了标注文件记得删掉 cache 重新生成否则用的还是旧标签。这个坑很隐蔽表现是训练 loss 正常下降但模型效果很差。4. 训练配置与超参数调整把模型训到收敛4.1 模型结构选择与预训练权重加载Yolov5 提供 s、m、l、x 四个尺度的模型参数量和精度递增。选哪个取决于你的数据量和部署硬件模型参数量适用场景yolov5s7.2M数据量小、边缘设备部署yolov5m21.2M平衡精度和速度yolov5l46.5M服务器端、追求精度yolov5x86.7M数据量大、精度优先我一般先用 yolov5s 跑通全流程确认数据和配置没问题后再换大模型。训练命令python train.py \ --data data/mydata.yaml \ --weights yolov5s.pt \ --img 640 \ --batch 16 \ --epochs 100 \ --device 0 \ --project runs/train \ --name mydata_exp1--weights yolov5s.pt表示从预训练权重开始微调这比从头训练收敛快得多小数据集上尤其重要。--img 640是输入分辨率如果你的目标很小可以提到 1280但显存占用会翻倍。--batch 16根据显存调整8G 显存跑 640 分辨率大概能到 1612G 能到 32。4.2 关键超参数的含义与调整方法训练相关的超参数分两部分命令行参数和 hyp yaml 文件。命令行控制训练流程hyp 控制优化器和损失权重。命令行里最需要关注的几个--epochs训练轮数。小数据集 100 到 300 轮大数据集 300 到 500 轮。看验证集 mAP 不再提升就可以停。--batch-size批大小。越大梯度越稳但显存占用高太小可能导致 BN 层统计不准。--lr0初始学习率默认 0.01。微调时建议降到 0.001 到 0.01 之间。--patience早停耐心值默认 100 轮验证集无提升就停。--freeze冻结层数小数据集可以冻结 backbone 前几层防止过拟合。hyp 文件里影响大的lr0: 0.01 lrf: 0.01 # 最终学习率 lr0 * lrf momentum: 0.937 weight_decay: 0.0005 warmup_epochs: 3.0 box: 0.05 # 框回归损失权重 cls: 0.5 # 分类损失权重 obj: 1.0 # 目标置信度损失权重如果发现模型框的位置不准调大box如果类别分错多调大cls。但别一次改太多每次只动一个参数观察验证集指标变化。4.3 训练过程监控与中断恢复训练启动后终端会实时打印每个 epoch 的 loss 和 mAP。同时runs/train/mydata_exp1/下会生成weights/best.pt验证集 mAP 最高的权重weights/last.pt最后一轮的权重results.csv每轮指标记录results.pngloss 和 mAP 曲线图判断训练是否正常看几个信号box_loss、obj_loss、cls_loss 三个都应该整体下降如果某个震荡剧烈或者不降说明对应部分有问题。mAP0.5 应该逐步上升到后期趋于平稳。如果训练中断了用--resume恢复python train.py --resume runs/train/mydata_exp1/weights/last.pt这会从上次中断的 epoch 继续优化器状态也会恢复。注意--resume后面跟的是 last.pt 的路径不是 best.pt。5. 避坑与排查那些让我熬夜的翻车现场5.1 显存溢出与 batch size 调整现象训练启动几秒后报CUDA out of memory或者跑到某个 batch 突然崩。原因显存不够。可能是 batch size 太大、输入分辨率太高、模型太大也可能是 mosaic 增强时四张图拼接导致单张实际分辨率翻倍。解决先把 batch size 减半试还不行就降--img到 416 或 320。如果用的是 yolov5l 或 x换 s 或 m。另外可以开--amp混合精度训练显存占用能降 30% 左右速度也更快。实在不够就上梯度累积--batch-size设小但--accumulate设大等效于大 batch。5.2 训练 loss 不下降或 mAP 为 0现象训练跑了十几轮loss 几乎不变验证集 mAP 一直是 0。原因最常见的是标注格式不对。比如坐标没归一化、class_id 从 1 开始而不是 0、txt 文件和图片文件名不对应。其次是 yaml 里nc和实际类别数不一致或者names顺序和标注 id 对不上。解决先用python utils/general.py里的可视化函数抽查几张训练图确认框画得对。再检查 labels 目录下 txt 文件数量是否和 images 一致。最后确认 yaml 里nc等于len(names)。如果都没问题把学习率降到 0.001 再试。5.3 验证集指标虚高但实际推理效果差现象训练日志里 mAP0.5 到了 0.9 以上但拿新图片推理时漏检严重。原因训练集和验证集分布太接近或者数据增强过强导致模型只学会了增强后的模式。也可能是验证集里重复图片太多相当于变相泄漏。解决重新划分数据集确保验证集里的场景、光照、目标尺度有足够多样性。降低 mosaic 和 mixup 的概率或者后期关闭增强。另外检查验证集有没有和训练集重复的图片用文件哈希去重。5.4 推理时类别名显示错误或框重叠现象detect.py 跑出来框的位置对但类别名是错的或者同一个目标出了好几个框。原因类别名错误通常是 yaml 里names顺序和训练时不一致。框重叠是 NMS 的 IoU 阈值设得太高。解决推理时用的 yaml 必须和训练时完全一致包括names的顺序。NMS 阈值通过--conf-thres和--iou-thres调整默认 0.25 和 0.45漏检多就降 conf框重叠多就降 iou。5.5 从 Yolov5 迁移到 Yolov8 时的格式差异现象Yolov5 训好的数据直接喂给 Yolov8 报错。原因Yolov8 的 yaml 格式和 Yolov5 略有不同主要是路径写法和 split 定义方式有变化。解决Yolov8 的 yaml 里train、val直接写图片目录路径不需要path字段。另外 Yolov8 默认不需要单独的 labels 目录它会在图片同级目录找同名 txt。迁移时把目录结构调整一下或者改 yaml 里的路径指向。6. 推理验证与模型导出让训练结果真正能用训练完成后第一件事是拿 best.pt 在测试集上跑一遍确认指标和训练日志一致python val.py \ --data data/mydata.yaml \ --weights runs/train/mydata_exp1/weights/best.pt \ --img 640 \ --task test--task test指定用测试集评估不指定默认用验证集。输出会打印每个类别的 precision、recall、mAP0.5、mAP0.5:0.95。如果测试集指标比验证集低很多说明模型过拟合了需要加数据或加正则。推理单张图片或整个目录python detect.py \ --weights runs/train/mydata_exp1/weights/best.pt \ --source test_images/ \ --img 640 \ --conf-thres 0.25 \ --iou-thres 0.45 \ --save-txt \ --save-conf--save-txt会把检测结果存成 YOLO 格式的 txt--save-conf会在 txt 里带上置信度。输出默认在runs/detect/exp/下。如果要做视频检测--source换成视频文件路径就行。实际部署时通常需要导出成 ONNX 或 TensorRTpython export.py \ --weights runs/train/mydata_exp1/weights/best.pt \ --include onnx \ --img 640 \ --batch 1导出的 onnx 文件可以用 onnxruntime 加载推理也可以进一步转 TensorRT 加速。注意导出时的--img要和训练时一致否则精度会掉。如果部署到边缘设备建议导出时加--dynamic支持动态 batch但部分推理引擎对动态 shape 支持不好需要实测。我自己的习惯是每次训完模型先跑 val 确认指标再跑 detect 看几张实际图片的检测效果最后导出 onnx 用 onnxruntime 跑一遍确认精度无损。这三步走完模型才算真正可用。希望帮到你。本文还有配套的精品资源点击获取
返回列表