YOLOv5+PyTorch训练自定义数据集：完整流程与避坑指南

简介面向目标检测初学者与需要快速搭建自定义模型的开发者这份YOLOv5PyTorch实战教程覆盖了从环境准备、数据标注格式整理、数据集划分到模型训练、指标评估与推理部署的完整链路。资源包共69个文件压缩后仅13.81MB以Python脚本、YAML配置、图像样本和说明文档为主囊括训练/测试代码、模型结构配置、示例图片与多份README目录划分清晰便于按需查阅。目前已有302人学习下载。借助包内的数据划分与边界框检查脚本可快速清洗并校验自定义数据集提供s/m/l/x多尺度模型配置方便根据显存与精度灵活选择另含Dockerfile、权重下载脚本等实用工具简化运行环境搭建。通过该资源读者能理解YOLO单阶段检测的核心原理掌握修改数据配置、启动训练、监控损失与精度、执行测试可视化的完整方法大幅降低从零训练自定义检测模型的入门门槛。1. 训练自己的YOLOv5Pytorch数据集多数人栽在第一步“目标检测-使用Yolov5Pytorch训练自己的数据集-超详细流程教程-优质项目实战.zip”这类压缩包在搜索里铺天盖地标题自带“适合0基础纯小白”但真正解压后照做你会发现没有一篇讲清楚环境为什么装不上、标注文件为什么读不到、mAP为什么是零。我按自己跑过多个检测项目的经验把这条链路重新走一遍环境搭建、数据集制作、参数调整、报错排查、结果验证。每个环节都给可直接复制的命令和脚本同时把最容易翻车的地方单独拎出来说明。无论你是要做鸟类目标检测的数据集还是先拿YOLOv5练手入门这篇流程都适用。2. YOLOv5环境搭建CUDA、PyTorch、依赖三步对齐先说结论YOLOv5对PyTorch版本很宽容但对CUDA和PyTorch的匹配关系非常敏感。我见过最典型的翻车是明明nvidia-smi显示CUDA 12.x装的PyTorch却只编译到cu118训练到某个卷积层直接抛no kernel image is available。这个报错和显卡坏了没关系本质是驱动版本低于PyTorch运行时需要的CUDA版本。2.1 先分清驱动版本和PyTorch运行时版本CUDA不对齐的真实原因先执行两条命令nvidia-smi nvcc -Vnvidia-smi右上角显示的CUDA Version是驱动最高支持的CUDA版本它是一个“上限”不代表系统里真正装了对应版本的CUDA Toolkit。nvcc -V显示的是命令行可用的CUDA Toolkit版本如果提示command not found说明你没装独立Toolkit。PyTorch通过pip安装的wheel自带CUDA运行时所以缺nvcc也能训练但遇到需要现场编译C扩展的依赖时缺Toolkit会少一条路。判断原则很简单驱动支持的CUDA版本要大于等于PyTorch编译时用的CUDA版本。比如驱动显示12.2装cu121、cu122甚至更低的wheel都能跑驱动只显示11.4就不要硬装cu118以上的PyTorch否则就会出前面说的kernel image报错。安装PyTorch前我建议先用conda隔离环境conda create -n yolo python3.8 -y conda activate yolo为什么推荐Python 3.8YOLOv5的依赖在3.8到3.10上踩坑最少太新的Python版本可能让numpy、opencv找不到现成wheelpip会现场编译Windows机器上常常因为缺VC编译链直接失败。如果机器上没装Anaconda用Miniforge也完全可以命令一样。2.2 安装依赖先装PyTorch再装requirements.txt很多人直接执行pip install -r requirements.txt想着一步到位。实际上requirements.txt里的torch版本可能和你选好的CUDA不匹配pip会帮你“随便”装一个能跑但跑不动GPU的版本。更稳的做法是先手动装PyTorch确认GPU可用后再装剩下的依赖pip install torch torchvision装完立刻验证python -c import torch; print(torch.__version__, torch.cuda.is_available())输出True不代表后面一定没坑但输出False说明PyTorch的CUDA装配不对。如果返回False需要到PyTorch官网的whl页面换成对应CUDA版本的安装命令重装。确认没问题后再执行pip install -r requirements.txt注意requirements.txt里的版本号是作者锁过的不要为了“升级”随手放宽numpy或opencv版本。YOLOv5的坐标变换、锚框计算对部分numpy行为有隐性依赖版本太新或太旧都可能让结果出现偏差。这种问题最难查因为它不报错只是框偏移。2.3 用detect.py跑一次官方示例给环境做体检环境装好后先用官方自带的测试图验证python detect.py --weights yolov5s.pt --source data/images --conf-thres 0.25第一次运行会自动下载yolov5s.pt权重。检测结果输出在runs/detect/exp/目录。data/images下默认有两张测试图一张bus一张zidane如果两张都能画出框且类别标得对环境就算稳了。此时别急着拿几千张图开训先用小数据集把整个流程跑通再上全量数据出问题也容易定位。跑通之后顺手看几行日志里面包含模型层数、参数量和推理耗时这些数据可以作为后续训练的基准。如果GPU显存占用一直很低说明PyTorch可能还在用CPU推理需要回到上一节重新匹配CUDA版本。这一步属于“后悔药”环节现在花十分钟确认比训练到一半再排查划算得多。3. 准备数据集从标注格式到YOLO txt的落地流程准备数据是最能拉开项目周期的环节。很多人以为标注完就结束了实际还要处理格式转换、训练验证划分、标签一致性三个问题。这里用具体例子走一遍假设你在做鸟类目标检测的数据集目标类别分为bird、nest、eggs三类。3.1 标注格式怎么选VOC XML、COCO JSON、YOLO txt对比标注阶段用哪个格式主要看标注工具。LabelImg默认输出VOC式XML新版本也支持直接存YOLO txtLabelMe输出JSON如果你下载的是CCPD车牌数据集标签又是txt格式但字段含义和YOLO不完全一致。实际经验是标注时统一用一种格式训练前全部转成YOLO txt不要让train.py去适配多种格式。格式文件组织坐标表达常见来源VOC XML一张图对应一个.xmlxmin/ymin/xmax/ymax 像素坐标LabelImg默认、VOC系列数据集COCO JSON整个数据集一个.json像素坐标加分割多边形COCO、大型标注工程YOLO txt一张图对应一个.txt归一化中心点坐标和宽高YOLOv5直接读取做遥感检测时手头数据像HRSC2016VOC格式居多要转YOLO txt转换脚本可以直接复用。之所以强调转成YOLO txt是因为YOLOv5数据加载器读它最快后续做mosaic等数据增强时坐标同步变换也最方便。3.2 把VOC标注转成YOLO格式转换脚本与四个边界坑假设你已经用LabelImg标了一批VOC XML目录结构如下dataset/ images/ bird_001.jpg bird_001.xml bird_002.jpg bird_002.xml labels/ # 待生成转换脚本import os import glob import xml.etree.ElementTree as ET # 按自己的类别清单修改顺序决定训练时的class_id classes [bird, nest, eggs] images_dir dataset/images labels_dir dataset/labels os.makedirs(labels_dir, exist_okTrue) for xml_path in glob.glob(os.path.join(images_dir, *.xml)): tree ET.parse(xml_path) root tree.getroot() size root.find(size) w int(size.find(width).text) h int(size.find(height).text) lines [] for obj in root.iter(object): name obj.find(name).text.strip() if name not in classes: continue cls_id classes.index(name) bnd obj.find(bndbox) xmin float(bnd.find(xmin).text) ymin float(bnd.find(ymin).text) xmax float(bnd.find(xmax).text) ymax float(bnd.find(ymax).text) # 像素坐标归一化到0-1范围 cx (xmin xmax) / 2 / w cy (ymin ymax) / 2 / h bw (xmax - xmin) / w bh (ymax - ymin) / h # 边界钳制防止标注越界 cx min(max(cx, 0.0), 1.0) cy min(max(cy, 0.0), 1.0) bw min(max(bw, 0.0), 1.0) bh min(max(bh, 0.0), 1.0) lines.append(f{cls_id} {cx:.6f} {cy:.6f} {bw:.6f} {bh:.6f}\n) if not lines: continue base os.path.splitext(os.path.basename(xml_path))[0] with open(os.path.join(labels_dir, base .txt), w, encodingutf-8) as f: f.writelines(lines) print(f{base}.txt - {len(lines)} objects)脚本逻辑不复杂解析XML里的size节点拿到原始宽高再把每个目标的bbox换算成归一化中心坐标和宽高。YOLO txt里不存像素坐标这行格式是数据加载器唯一认的东西。后面标注坐标系错了训练会悄悄把检测框学歪这也是为什么转换脚本里要同时做边界钳制。这里要提醒四个边界坑全是实际踩过的XML里的filename节点可能和真实文件名对不上。标注完再改名、批量重排是常事所以脚本要以XML文件名写输出txt不要用filename节点去映射图片。类别名大小写必须严格一致。Bird和bird是两回事标错一个这一批目标直接变背景mAP计算也会跟着崩。bbox越界很常见。手动标注时xmax大于图像宽度的情况不少脚本做了钳制但如果xmin大于xmax钳制救不了转换前应该先输出警告。标签文件为空时不要生成txt。一张图没有任何有效目标要么删掉这张训练图要么检查类别清单让空文件混进训练集只会干扰损失计算。3.3 训练集与验证集划分每类准备多少张才不容易翻车转换完成后按8:2或9:1划分训练和验证图片常见做法是用两个目录分别存图像dataset/ images/ train/ bird_001.jpg bird_002.jpg val/ bird_020.jpg labels/ train/ bird_001.txt val/ bird_020.txtYOLOv5有个隐式规则标签目录和图像目录平级出现把路径里的images替换成labels即可。也就是说images/train对应的标签目录必须叫labels/train不要叫label/train或train_labels。刚上手时最容易在这里踩坑路径对不上日志会提示找不到标签但不会报具体哪个路径错。每类目标最少准备多少张我的经验是300张上下比较稳100张也能训但mAP波动会很大。如果数据量不足优先加数据增强而不是直接上大模型。先拿一个小验证集跑一轮确认loss能降下来再全量训练。这里提到的训练集和验证集划分方式同样适用于CCPD这类现成数据集只是不需要标注直接按目录复制图片即可。4. 修改配置文件和训练参数从yaml到train.py首轮训练环境没问题、数据转换完了接下来就是最容易“白跑”的环节配置文件。很多训练失败案例并不是模型不行而是类别名没写对、路径没对齐、参数选得不合理。4.1 先改data/custom.yaml类别、路径和nc的关系训练前先写数据yaml放在仓库的data目录下# data/custom.yaml train: ../dataset/images/train val: ../dataset/images/val nc: 3 names: 0: bird 1: nest 2: eggstrain和val分别指向图像目录相对路径是相对于yaml文件所在位置解析。nc必须和names列表长度一致。这里有个关键点如果你用预训练权重模型yaml里的nc并不是最终决定因素data yaml的nc会在训练时覆盖模型输出层维度。names顺序会写入训练日志和结果输出必须和转换脚本里的classes顺序完全一致否则类别ID对不上训练结果就是“能收敛但检测全错”。有些人还会去改models/yolov5s.yaml里的nc这个改动只在从头训练时有意义。如果你用--weights yolov5s.pt做迁移学习data yaml里的nc会自动重建输出层模型yaml改不改不影响。新手没必要两边都改改一处即可。4.2 train.py必调参数batch-size、img、epochs和显存上限第一次训练建议用官方预训练权重起步优先选yolov5s。如果你处理的是小而密的目标再考虑yolov5m或直接调大输入分辨率。我一般用的命令python train.py \ --data data/custom.yaml \ --weights yolov5s.pt \ --img 640 \ --batch-size 16 \ --epochs 100 \ --cache几个核心参数的含义和建议参数默认值我常用的调整说明--img640640或1280输入分辨率小目标提升明显但显存开销近似翻4倍--batch-size16按显存减半出现OOM时优先降这个参数--epochs10030试跑先跑30轮看曲线收敛趋势好再放量--weightsyolov5s.pt预训练权重想从零训练时传空字符串--cacheFalse显存换内存提升数据读取速度但内存小于32G建议别开--patience10030早停耐心值验证指标长期不升自动停止如果你处理的是移动小目标检测这类问题比如无人机俯拍图里的行人或车辆--img 1280带来的收益通常比换更大的模型更明显代价是显存需求成倍增长。显存不够时先把batch-size减半再把img降到640不要两个参数同时硬扛。刚开始训练时不要直接跑100轮先跑30轮看损失曲线。如果你发现box loss在下降但很慢可以检查标注框是否准确如果loss震荡不收敛大概率学习率不合适。YOLOv5默认超参数在公共数据集上表现不差但你的数据分布和COCO差别大时需要单独调。4.3 训练日志怎么看loss、P/R、mAP和混淆矩阵训练启动后日志会逐轮输出Epoch gpu_mem box obj cls labels img_size 1/100 4.62G 0.08321 0.02904 0 17 640这里box是边框回归损失obj是目标置信度损失cls是分类损失。如果cls一直为0先不要慌第1个epoch分类分支本来就可能还没激活等跑到第10轮再判断。每轮结束runs/train/exp/下会生成results.png里面有box loss、obj loss、cls loss、precision、recall、mAP曲线。看曲线时我一般抓三个特征三条loss线整体向下且mAP稳步上升属于正常收敛train loss一直降但mAP不升反降是过拟合信号考虑加数据增强或提前停止loss线横跳不下降多半是学习率过大或标签噪声太大。混淆矩阵图能直接看出哪些类别互相误检比如bird被识别成nest这种问题从曲线里看不出来必须看混淆矩阵才能定位。5. YOLOv5训练常见问题避坑现象、原因、解决一条龙训练期总会遇到几种反复出现的报错。这里挑出实际项目里高频的问题按现象、原因、解决三步写清楚。每次都值得把触发条件记录下来下一次换新数据集直接绕开。5.1 CUDA out of memory先别急着换卡调两个参数现象训练开始后第一个batch或中途某个epoch报RuntimeError: CUDA out of memory有时提示“Tried to allocate 2.00 GiB”。原因模型大小、batch size、输入分辨率三者乘积超过显存容量。PyTorch在训练时还会临时缓存中间激活值不是只看模型权重占多少显存。解决先把--img缩到512再看显存占用还不够就把batch-size降到4或2。batch降到4还溢出检查是不是浏览器、显示进程占了大量显存。还有个隐藏坑YOLOv5默认开多尺度训练每10个epoch随机变换输入尺寸某个epoch忽然变大就可能爆显存。如果机器显存紧张手动固定img更保险。5.2 No labels found标签路径和图像目录没对齐现象训练日志显示能正常启动但提示No labels found in .../images/train随后大量图像被跳过训练几乎没进展。原因yaml里train指向图像目录但找不到与之配对的labels目录或者txt文件全是空的。解决先确认images/train上级是否存在labels/train再确认txt文件名和jpg文件名完全一致最后打开一个txt看第一行的class_id是否超过nc减1。YOLOv5只认images和labels这两个目录名配对关系你把目录改成imgs或label它不会自动识别。路径层次越深越容易出错按上面3.3的结构逐层核对比反复改yaml更省时间。5.3 损失变成NaN坏数据和过大学习率都是祸根现象训练日志里的loss突然变成NaN之后跑再多轮也不恢复。原因学习率太高导致梯度爆炸或者训练图片中有全黑、全白、损坏的坏图某些标签坐标出现inf值。这种问题不报路径错误排查起来最费时间。解决先写脚本扫描所有txt标签查找包含inf或nan的行以及尺寸为零的bbox。如果转换脚本做过钳制问题多半出在原始XML里size节点缺失。超参数层面换用低学习率的hyp文件比如仓库里自带的hyp.scratch-low.yaml把初始学习率降一档。处理完坏数据再训练NaN通常就消失了。这种“看不出原因”的问题很多人喜欢归为玄学实际上八成是数据文件问题。5.4 mAP一直为0类别ID和names列表不匹配是头号嫌疑现象训练能跑完曲线看起来也正常但验证mAP等于0检测框一个都出不来。原因最常见的是names列表顺序和转换脚本里classes顺序不一致。比如训练时0号是bird转换时0号却是eggs标签匹配全乱。另一种可能是val目录里混进了没有标签的图片拉低了评估结果。解决取一张验证集图片连同txt标签打印出来再指定best.pt单独检测这一张图比对输出类别名。如果类别ID和检测框都对还找不到问题就把--conf-thres降到0.05试一次。小目标数据集的置信度普遍偏低阈值设太高会导致看起来“一个都没检测到”。这也是新手最容易判断错的方向以为是模型没学好其实是评估和推理时的置信度口径不一致。5.5 导出ONNX报错opset与onnx版本的老搭配问题现象执行python export.py --weights best.pt --include onnx时在torch.onnx.export阶段报错或者导出成功后用ONNX Runtime推理结果不对。原因torch、onnx、onnxruntime三者的版本组合没有对齐。老权重模型导出到新版本时某些算子不兼容也很常见。解决显式指定--opset 11让导出和推理使用同一套算子集不要依赖默认值。同时保持requirements.txt里onnx相关版本不变不要手贱升级。如果你想部署到RK3568这类边缘设备工具链通常以ONNX或tflite为输入先把导出跑通再考虑量化和精度损失的问题。导出后用onnxruntime加载模型喂一个随机张量跑一次能输出正常形状的结果再进工具链这样能隔离问题阶段。6. 用best.pt对陌生图片做回归验证再决定下一步调参训练完成后真正让模型“能用”的标志不是看训练曲线多漂亮而是拿一批它没见过的图片按真实使用场景跑一遍。我习惯把训练时没用过的测试图片单独放进一个目录然后执行python detect.py --weights runs/train/exp/weights/best.pt --source test_images/ --conf-thres 0.25跑完我会重点判断两件事漏检多不多误检多不多。漏检多就把置信度从0.25降到0.1再看检测框的置信度分布误检多就把阈值往上调。模型被用在什么场景就用什么场景的照片做验证不要拿网上的公开图随便顶替分布差异会掩盖真实性能。如果训练时用1280分辨率检测时也要保持一致否则精度会明显下降。如果项目需要落地部署我下一步会导出ONNX作为转换到边缘设备工具的中间格式python export.py --weights runs/train/exp/weights/best.pt --include onnx --opset 11导出成功后用ONNX Runtime做一次快速验证import onnxruntime as ort import numpy as np sess ort.InferenceSession(best.onnx) input_name sess.get_inputs()[0].name dummy np.random.randn(1, 3, 640, 640).astype(np.float32) outputs sess.run(None, {input_name: dummy}) for o in outputs: print(o.shape)YOLOv5的ONNX输出是多尺度特征图形状一般是1x(5类别数)*3xHxW。能看到几个尺度的输出说明导出基本成功。部署阶段还需要写NMS后处理或者用ONNX Runtime自带的NMS算子这部分和训练是两个独立问题。顺带说一句之后如果你从YOLOv5切到YOLOv11系列数据yaml和标签格式还能复用底子是一样的但环境配置和部分参数名会变。最后说一个影响我很久的习惯每次训练结束我会把best.pt、results.png、验证集检测结果和当次使用的data yaml、完整训练命令一起备份到按日期命名的目录。这样过两周再回来调参不用靠记忆猜上次改了哪些参数。尤其是换数据集、改类别名这类操作最怕的就是“优化过但忘了改了什么”。这个习惯帮我省下大量返工时间希望帮到你。本文还有配套的精品资源点击获取