YOLOv5垃圾分类识别实战:从数据标注到TensorRT部署全攻略

发布时间:2026/10/11 17:14:05
YOLOv5垃圾分类识别实战:从数据标注到TensorRT部署全攻略 简介基于YOLOv5的垃圾分类识别检测项目面向深度学习初学者与计算机视觉实践者解决居民生活垃圾自动分类与类别定位问题。压缩包内共包含九千余个文件其中以八千多张jpg图像、六百余个txt标注文件、近两百个xml标签以及三十个yaml配置文件为主另附若干Python脚本、权重文件与Docker部署文件整体约五百兆字节目录结构清晰可完整支撑从数据准备到模型推理的各个环节。项目清晰展示了数据集格式转换、探索性数据分析、环境安装、YOLOv5源码修改以支持中文标签、训练集与测试集自动划分、配置文件调整、WeightsBiases训练可视化、网络模型训练以及最终测试与性能统计的完整实验流程。已有两千二百余人学习下载适合希望复现垃圾分类检测全流程、理解YOLOv5工程化落地方式以及参考目标检测项目完整目录组织与调参思路的读者。1. 为什么我用 YOLOv5 做垃圾分类识别一次框出几十件垃圾的底气垃圾分类这事做图像分类方案的人会有个根深蒂固的惯性——觉得把图片扔进一个分类网络输出一个标签就算完事。但真正上手去小区垃圾投放点拍一段视频就会明白镜头前往往是多个饮料瓶、纸箱、快递袋堆叠在一起分类网络只能给整张图一个类别完全不够用。YOLOv5 这类单阶段目标检测模型一次推理就能同时给出几十个检测框和每个框对应的垃圾类别这才是从「图像分类」走向「场景落地」的关键一步。本篇笔记面向正在做毕业设计、课程设计或者准备把垃圾分类接入边缘设备的开发者我把模型选型、数据集转换、训练参数、部署踩坑这整条链路的经验完整拆给你。2. YOLOv5 检测原理与选型理由为什么是 YOLOv5 而不是 YOLOv82.1 从 YOLOv3 到 YOLOv5anchor 机制与损失函数的变化逻辑YOLOv5 的检测头延续了 YOLOv3 的 anchor-based 设计这可能是 2025 年还选择它作为毕设主模型最扎实的理由它是一个「稳定、可复现、社区案例最多」的方案。网络结构上输入图像先经过 CSPDarknet 骨干网络做特征提取再用 PANet 结构把深层语义信息和浅层纹理信息做融合最后在三个不同尺度的特征图上做预测。小尺度特征图感受野大适合检测纸箱这类大物件大尺度特征图对细节更敏感适合捡出纽扣电池、烟头这类小目标。训练阶段的损失函数是一个多任务加权和定位损失GIoU Loss、置信度损失BCE With Logits、分类损失BCE With Logits。我在调参时最常看的指标就数这三部分的变化趋势——定位损失收敛不理想最常见的原因是标注框本身画歪了置信度损失下不去优先检查数据里是否存在大量前景背景比例失衡的样本。2.2 模型规模选择n/s/m/l/x 五种尺寸的取舍YOLOv5 官方仓库提供 n、s、m、l、x 五种预先设计好的模型正好对应从边缘设备到服务器端的算力梯度。我的建议是毕业设计先跑yolov5s它显存占用大约 4GB能够在 1080Ti 甚至 1650S 这种中端显卡上训练推理速度也能到 60FPS。如果你最终要部署到 Jetson Nano 或者树莓派这类低算力设备直接选yolov5n它比 s 轻约 40%mAP 只掉 3-5 个点差异主要体现在硬币、瓶盖这类小目标上。经验之谈不要在入门阶段直接上yolov5x。垃圾类别之间边界模糊的不少比如「果皮」和「湿纸巾」外观可能相近大模型对数据量非常贪婪几千张图片训 x 往往得不到预期的涨点反而容易把训练时间拉长两三倍。2.3 环境搭建实操虚拟环境与 CUDA 版本对齐我一般会靠 conda 把环境隔离开避免不同项目之间包冲突conda create -n yolo_garbage python3.8 conda activate yolo_garbage pip install torch1.13.1 torchvision0.14.1 --index-url https://download.pytorch.org/whl/cu117 git clone https://github.com/ultralytics/yolov5.git cd yolov5 pip install -r requirements.txt这里有几个关键点需要说明先装 PyTorch 再装 requirements.txt可避免 requirements 把默认的 CPU 版 torch 装上。cu117表示 CUDA 11.7如果你的显卡驱动只支持 CUDA 11.3就把cu117换成cu113。Net 20 系列显卡直接从cu117起步老显卡才需要降版本。3. 垃圾数据集准备与格式转换从 VOC XML 到 YOLO TXT 的完整流程3.1 数据获取策略公开数据集与场景自采怎么配合垃圾分类数据集现在有不少公开选项常见的有某高校发布的鸿鹄数据集、某公司开放的垃圾检测集合。但我的实际经验是直接拿来训练mAP 可能很漂亮一到学校、公司楼下实测就翻车——公开数据集的拍摄场景、光照条件、垃圾桶类型和你落地的环境差异太大。正确做法是拿公开集做预训练再从你的目标场景里补充拍摄 300-500 张现场照片做微调。拍摄环节有四个硬性要求垃圾桶要放在画面中心偏下垃圾至少要覆盖画面 20% 以上面积光线要覆盖白天、傍晚、灯光三个场景每个类别至少拍 20 个不同个体避免模型只认出某一种特定饮品瓶。3.2 使用 LabelImg 标注类别一致性比标注速度更重要标注工具选labelImg即可pip 安装后启动pip install labelImg labelImg images classes.txt软件操作层面不多说重点提醒类别标签顺序classes.txt里的类别排列顺序必须和后面训练配置里的names列表严格一致。例如第 0 行是glass_bottle训练配置里names第 0 项也必须是glass_bottle多了少了一列都会导致标签错位行为表现是训练时 loss 正常下降推理时检测框永远给出错误的类别。3.3 VOC 格式转 YOLO 格式坐标归一化转换脚本LabelImg 默认保存的是 PASCAL VOC 格式即xml文件。YOLOv5 训练需要的是 txt 格式每行对应一个目标格式为class_id x_center y_center width height其中 x_center、y_center、width、height 都归一化到 0-1。转换脚本是这种任务里最不需要「重新发明轮子」的部分直接用我调好的版本改路径import xml.etree.ElementTree as ET import os classes [glass_bottle, plastic_bottle, paper_box, cigarette_end] def convert_xml_to_txt(xml_path, output_dir, image_width, image_height): tree ET.parse(xml_path) root tree.getroot() lines [] for obj in root.iter(object): class_name obj.find(name).text if class_name not in classes: continue class_id classes.index(class_name) box obj.find(bndbox) xmin float(box.find(xmin).text) ymin float(box.find(ymin).text) xmax float(box.find(xmax).text) ymax float(box.find(ymax).text) x_center (xmin xmax) / 2 / image_width y_center (ymin ymax) / 2 / image_height width (xmax - xmin) / image_width height (ymax - ymin) / image_height lines.append(f{class_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}) output_path os.path.join(output_dir, os.path.splitext(os.path.basename(xml_path))[0] .txt) with open(output_path, w) as f: f.write(\n.join(lines)) # 示例调用 convert_xml_to_txt(data/annotations/001.xml, data/labels, 1920, 1080)脚本本身不做图像尺寸读取因为每个 xml 的 size 节点里其实存着宽高生产级脚本应从 xml 里解析这个值。我这里为了简化将它作为函数参数传入你实际用的时候解析root.find(size/width)比手动指定更稳妥。类别顺序classes必须与标注时的classes.txt一致写成列表后便于多次复用。3.4 数据划分按文件夹而不是按随机种子划分训练集、验证集、测试集的划分YOLOv5 的train.py接受--train和--val参数指向存放图片的文件夹路径。注意它并不接受一个txt文件列表那是 YOLOv3 时代的方式。所以最简单有效的划分方法是物理分目录images/train、images/val、labels/train、labels/val。我习惯用random.sample做索引洗牌并固定随机种子保证后续消融实验对比公平。比例上场景自采数据少可以按 7:2:1 切分。另外要强调一个问题如果你既做训练又要评测最终模型务必单独切出images/test目录train.py 只在训练时看到 train 和 valtest 留到test.py阶段使用。很多初学者把 val 当 test 用导致最后报告的 mAP 有信息泄露的成分。4. 基于 YOLOv5 训练垃圾分类模型核心参数与调优实操4.1 训练启动命令从预训练权重开始而不是从零开始直接用 ImageNet 预训练的权重做迁移学习是目标检测任务里提升收敛速度最有效的手段。YOLOv5 的--weights参数同时扮演了两个角色骨干网络加载预训练参数检测头则根据你的类别数重新初始化。这意味着即使你的垃圾类别与 COCO 的 80 类完全不同预训练依然有效——特征提取层学的是通用的边缘、纹理、形状感知能力。python train.py \ --data garbage.yaml \ --weights yolov5s.pt \ --img 640 \ --batch 16 \ --epochs 150 \ --workers 8 \ --device 0关键参数逐个说明。--img 640是训练输入尺寸分辨率越高小目标信息保留越完整但显存消耗按平方增长。--batch 16是总 batch size——如果是单卡batch 和显存直接相关8GB 显存下--img 640 --batch 16是安全阈值。--epochs 150面对几千张垃圾分类图片够用再多容易过拟合。--workers 8是数据加载线程数Windows 上设 0 或 2 可以规避 DataLoader 卡死问题Linux 直接 8。garbage.yaml是数据配置内容结构如下path: ./garbage_dataset train: images/train val: images/val test: images/test nc: 4 names: [glass_bottle, plastic_bottle, paper_box, cigarette_end]4.2 训练过程中的 loss 曲线判读什么信号值得停训YOLOv5 训练结束后会在runs/train/exp目录下生成results.png它绘制了train_loss、val_loss、mAP0.5、mAP0.5:0.95等六条曲线。我判断模型是否训练到位有三个信号第一train_loss最终稳步下降到初始值的 30% 上下并在最后 20 个 epoch 不再出现剧烈抖动。第二val_loss呈现「先下降、后平台、再上升」的典型过程最低点就是早停时机后续训练只会过拟合。第三mAP0.5普遍以 95% 为分水岭——垃圾分类任务里如果类别状态单一比如只有干净背景和待扔垃圾mAP0.5 到 97% 以上才算合格。4.3 超参数调优anchor 尺寸与图片尺寸的联动YOLOv5 在训练时会自动重新聚类 anchor你不需要手工指定但--img尺寸直接决定了 anchor 尺度的分布。具体来说--img 640训练的模型推理时使用--img 1600往往导致小目标检测性能反而下降因为 anchor 尺度是为 640 分辨率设计的。推理尺寸应与训练尺寸保持一致最多上下浮动 30%。如果你数据集中同时存在「整张外卖包装袋」「烟头」这种极端尺度差异我建议打开hyp.scratch-low.yaml修改anchors相关参数至少在anchor聚类时把kmeans_iters从默认 1000 提高到 5000避免聚类陷入局部最优。这个参数在utils/autoanchor.py中定义直接改该文件里的常量即可。4.4 类别不均衡的损失函数修正垃圾分类数据集里「塑料瓶」的样本数量可能是「药品」的二十倍。模型会偏向多数类导致少数类 mAP 极低。YOLOv5 的损失函数为每个类别分配独立权重做法是修改hyp.scratch.yaml中的cls_pw参数。如果你的类别只有 4 个可以通过--cls参数设置类别权重列表例如给样本少的类别更高的损失权重。实操中我常用的另一种方式是过采样把样本量低于 100 张的类别在数据集中复制两到三份。复制时注意不要简单将一个文件重复放置而应该做数据增强副本比如翻转、旋转 30 度、加高斯噪声这样可以让数据增强真正发挥作用而不会让模型见过多的同一张原始图片。5. 垃圾分类训练与部署避坑六个我用数据换来的教训5.1 显存溢出 OOM调低 batch 还不够问题在图像尺寸现象batch 降到 4 依然在训练十几个 step 后报 CUDA out of memory。原因YOLOv5 在训练时除了输入图像还会为每个图像的 anchor 输出计算损失显存峰值出现在 forward-backward 交界处。--img 1280的实际显存需求是--img 640的 4 倍回不去了是真的回不去了。解决先确认--img不是 1280其次给train.py增加--cache选项把数据预加载到内存中减少 I/O 等待时间。如果显存仍然不够修改--batch -1让 YOLOv5 自动根据显存计算能用的最大 batch。这才是最不折腾的方案。5.2 迁移学习后模型对真实场景失效公开数据集和实际环境差太远现象训练时 mAP0.5 到了 98%拿到学校走廊垃圾桶实测检测框来回抖类别误判率高。原因公开数据集的图基本是俯视拍摄、单一物品、干净背景而实际地面角度是斜视、多物品堆叠、各种杂乱的背景纹理。解决把实测场景的照片加入训练集并联合标注。最低门槛是补拍 200 张真实场景图加入训练集并重新标注。数据增强增强不过「领域鸿沟」这是特征分布层面的差异。5.3 推理时某些类别永远检测不出anchor 尺度与类别目标太不匹配现象纸箱和塑料瓶都能框出来烟头一个都检不出甚至没有低置信度的置信度输出。原因烟头在 640x640 分辨率下可能只有 30x15 像素属于典型小目标。YOLOv5 默认的 anchor 最小尺度是 4x4stride 32 下的最小框但这个尺度在垃圾场景里不够丰富。解决让 autoanchor 重新聚类。在训练日志里找到anchors输出对比是否覆盖了小目标尺度。如果聚类结果最小 anchor 仍在 50x50 附近检查标注框是否把烟头的 bbox 画得太宽松将 bndbox 从「烟头外围一圈」收窄到「烟头本身」anchor 聚类会跟着调整。5.4 ONNX 导出失败版本号错位带来的玄学报错现象执行export.py --weights best.pt --include onnx报错提示torch.onnx.export的simplify过程 KeyError。原因YOLOv5 不同 release 之间的export.py与 onnx-simplifier 版本不是一个配套集合pip 自动升级导致 API 断档。解决固定依赖版本。我的做法是执行pip install onnx1.13.1 onnx-simplifier0.3.10再配合torch 1.13.1这一组合兼容性最稳。输出 ONNX 后用onnxruntime跑一次收益对比确认与 PyTorch 推理结果一致再部署否则问题会从模型前移到推理框架。5.5 推理边界震荡NMS 阈值与暂态抑制参数的耦合现象视频检测时同一个瓶子的检测框每帧之间抖动严重甚至框的位置在瓶子和纸箱之间来回跳动。原因b区域重叠度高--conf-thres阈值设太低比如 0.1NMS 对低置信度框无法有效抑制。解决四个参数合理取值——conf_thres0.35iou_thres0.45max_det300。如果画面中垃圾堆叠严重可以适当上调iou_thres到 0.5但对模型输出不那么敏感的目标先试iou_thres0.45保持默认。视频场景最稳妥的方式是保留conf_thres0.4上下的配置。5.6 JPEG 压缩导致的标注框漂移现象用手机拍的照片经过微信传输后训练时 loss 下降很快但 mAP0.5 恢复正常到 90% 后无法再上升。原因手机原始图为 3000x4000 像素微信传输自动压缩到 800 像素标注是在原图上做的坐标在压缩后被隐式缩放而异常高压缩率产生的人工噪声影响了特征提取。解决训练前先统一所有图片到同一分辨率范围我偏好最长边 1280 并保持长宽比再执行一次标注坐标重算。这个过程交给脚本完成不要手工去改人眼处理 200 张以上图片的坐标校正一定会出错。6. 把训练好的模型跑在边缘设备上ONNX 转 TensorRT 半精度部署的完整流程用 YOLOv5 做的垃圾分类模型最终目标多半是跑在一块小计算板上。Jetson Nano、树莓派或者盒子设备它们的 CPU 跑 PyTorch 推理框架力不从心。第一步是导出 ONNX这在上一章避坑里已经提过。这里说从 ONNX 到 TensorRT 引擎的一条龙操作。也许你已经听过 TensorRT 这个词它本质上是 NVIDIA 出的推理加速引擎能把模型里的算子做层融合和精度校准。半精度 FP16 部署推理时间通常能压缩到 PyTorch 的 1/3 甚至更少显存占用也大幅下降。这也是 YOLOv5 落地的标准路径。导出的命令如下关键是搞清楚每个参数的意义python export.py --weights runs/train/exp/weights/best.pt \ --include engine \ --imgsz 640 \ --half \ --device 0--half开启半精度--include engine表示直接导出 TensorRT 引擎文件。这里可能有人会问为什么不先导出 ONNX 再转 TensorRT因为 export.py 实际上会自动完成这个流程PyTorch → ONNX → TensorRT。如果导出失败按上一章的版本对齐方案逐一排查。--imgsz必须与训练尺寸保持一致比如训练是 640这里就不要随意改成 320结果会变成无效的模型。之后用 TensorRT 推理python detect.py --weights runs/train/exp/weights/best.engine \ --source data/images \ --half \ --conf-thres 0.35 \ --iou-thres 0.45 \ --device 0.engine文件不是跨设备通用的。TensorRT 引擎和 GPU 架构强绑定在一块设备上生成的 engine 换到另一块就失效。我看到很多人踩这个坑后怀疑自己导出错了实际上这是正常的——你把best.engine和best.onnx一起拷贝目标设备上先用 ONNX 做精度对口再在设备本地重新生成 engine 保速度。部署时还需要对输入图像做预处理对齐YOLOv5 采用 letterbox 方式把不等比图像缩放至正方形像素值归一化到 0-1非 0-255推理输出经过 decode 还原到原图坐标后需要再做一次 letterbox 的逆变换才能画框。这个前处理只在你用自己的引擎接摄像头时才会遇到用 detect.py 会自动处理。用得多了以后我自己总结出的经验是每次训练完先强制走一遍「best.pt 导出 ONNX → 用测试集对精度 → 再用 ats 转化 TensorRT → 在目标设备上用一张真实场景图验收」这四步把前处理差异卡在模型集成初期。这个小习惯能省下后期排查部署问题的半天时间。希望帮到你。本文还有配套的精品资源点击获取

关于本文作者

来自尧图内容编辑团队

尧图内容编辑团队 内容团队

尧图内容编辑团队

本文由尧图网络内容编辑团队执笔。团队由资深项目经理、前端工程师与设计师组成,所有内容均来自亲手交付的真实项目,先讲清问题、再给出可落地的解法。尧图深耕北京网站建设十年,服务过京华建材集团、智造科技等各行业客户,把一线经验沉淀为可复用的行业观察。

  • 十年建站经验,覆盖建材、制造、服务、文创等
  • 项目经理把关选题与事实准确性
  • 工程师与设计师联合撰写专业细节
  • 统一编辑规范,保证文风与排版一致
  • 每月复盘转化数据,迭代选题方向

延伸阅读

相关资讯与近期热门内容

深度阅读推荐

建站决策前值得细读的三篇

网站改版的5个关键决策
2024-08-12

网站改版的5个关键决策

什么时候该改版、改到什么程度、如何避免流量掉光,京华建材集团改版复盘给出答案。

获取专属建站方案

看完文章,把您的行业与预算告诉我们,免费获取一份量身定制的官网建设方案与报价。

立即免费咨询