基于轻量化YOLOv5的生活垃圾分类方法及APP技术

本发明专利技术提供了一种基于轻量化YOLOv5的生活垃圾分类方法及APP，面向移动端设备手机设计部署，以原始的目标检测网络YOLOv5n为基线网络，采用引入ShuffleNetv2中的shuffle_block、GhostNet的Ghost模块、引入注意力机制SE以及知识蒸馏等技术手段，结合网络结构设计和知识蒸馏对基线网络实现轻量化，轻量化后的YOLOv5n模型相比原始的YOLOv5n模型，在很大程度上减少了模型的参数量和计算量以及提升了检测精度，将轻量化后的YOLOv5模型部署至移动端的生活垃圾分类APP中，可以实现对多种生活垃圾智能分类。垃圾智能分类。垃圾智能分类。

【技术实现步骤摘要】
基于轻量化YOLOv5的生活垃圾分类方法及APP


[0001]本专利技术属于计算机视觉
，具体涉及一种基于轻量化YOLOv5的生活垃圾分类方法及APP。

技术介绍

[0002]近年来，计算机视觉一直是科学界与工业界的研究热门点，作为计算机视觉中的基本视觉识别问题，目标检测在过去的几十年中得到了广泛地研究。目标检测是对图像或视频中的目标进行分类和定位的任务。目前，目标检测方法主要分为两类，分别是两阶段目标检测方法和一阶段目标检测方法。两阶段目标检测方法，例如FasterR
‑
CNN和CascadeR
‑
CNN和DetectoRS等。一阶段目标检测方法典型代表是YOLO系列以及SSD。两类目标检测方法各有优势：两阶段算法相对而言精度更高，尤其体现在定位上，而一阶段算法的速度普遍更快，更容易满足实际应用场景中的实时性需求。
[0003]目前人工智能技术应用的另一个重要趋势是在端侧平台上部署高性能的神经网络模型并能在真实场景中实时运行，如移动端/嵌入式端设备。然而现有的这些端侧平台的特点是内存资源少，处理器性能不高，功耗受限，这使得目前精度最高的模型由于对内存和计算资源的超额要求，使得其根本无法在现有的端侧平台上面部署且达到实时性的要求。为了使模型落地应用，需要设计轻量型网络结构或者通过知识蒸馏（DeepCompression）、剪枝(NetWorkPruning)、量化(ModelQuantization)等技术减少模型参数、计算量。大多数的模型压缩方法是在分类模型上进行的。常见的分类模型通常输入分辨率小、模型结构简单，而目标检测模型通常有着更大的输入分辨率，同时网络结构也更加复杂。目标检测模型可能还包括多个输入，以及一些后处理过程，这些都对模型压缩算法的设计提出了更高的要求。如果要在移动端部署目标检测网络，通常需要更浅的模型深度，更少的卷积核，但是不够深的网络模型往往在特征提取上不充分。因此一个平衡精度，检测速度和模型复杂度的目标检测网络显得至关重要。
[0004]目前生活垃圾分类主要面临以下问题：一是垃圾分类还是以人工分拣为主，因为人主观判断的不同造成分类错误从而导致无法从源头实现垃圾分类；二是智能分类垃圾桶尚未普及，而且现有智能垃圾分类系统一般采用的模型算法无法同时对多种垃圾进行分类，效率低下。此外，目前基于深度学习的生活垃圾分类方法针对的垃圾类别较少。如果提出一种方法能够快速便捷的区分日常生活垃圾属于可回收物、有害垃圾、厨余垃圾还是其他垃圾，就能够很大提高生活垃圾分类效率以及积极性。

技术实现思路

[0005]本专利技术的目的在于针对上述现有技术中的不足，提供一种基于轻量化YOLOv5的生活垃圾分类方法及APP，将轻量化目标检测网络在手机部署以实现日常生活垃圾的分类。
[0006]本专利技术采用以下技术方案实现：一种基于轻量化YOLOv5的生活垃圾分类方法，包括下列步骤：
S1、对生活垃圾数据集进行预处理，并将其划分为训练集、验证集和测试集；S2、将原始的YOLOv5n目标检测模型的Backbone替换为ShuffleNetv2的Backbone，得到初步轻量化的YOLOv5n模型；S3、使用COCO数据集对步骤S2得到的YOLOv5n模型进行训练得到训练权重best1.pt；S4、使用步骤S1选取的训练集和验证集对原始的YOLOv5s进行训练得到训练权重best2.pt，然后将best2.pt作为教师模型并使用步骤S3得到的best1.pt作为预训练权重对步骤S2得到的YOLOv5n模型进行知识蒸馏后得到训练权重best3.pt；S5、对步骤S2得到的YOLOv5n模型引入GhostNet中的Ghost模块，进一步得到轻量化的YOLOv5n模型并使用步骤S1选取的训练集和验证集对该YOLOv5n模型进行训练得到训练权重best4.pt，预训练权重为步骤S4得到的best3.pt；S6、对步骤S5得到的YOLOv5n模型引入SENet中的SELayer注意力模块得到最终轻量化的YOLOv5n模型，并使用步骤S1得到的训练集和验证集对该YOLOv5n模型进行训练得到训练权重best.pt，预训练权重为步骤S5得到的best4.pt；S7、将步骤S6的训练权重best.pt作为预训练权重，并以步骤S1选取的测试集作为步骤S6得到的最终轻量化的YOLOv5n模型的输入图像进行检测，最终检测完毕得到的图像中标明垃圾分类。
[0007]进一步地，所述步骤S1中包括：S101、使用LabelImg标注软件对生活垃圾数据集进行标签检测框调整；S102、将步骤S101调整好的数据集中的xml标签文件转换为YOLOv5可以使用的包含类别、中心点坐标、检测框高和宽的txt标签文件，其中txt标签文件中类别的编号需要连续，此外需手动修正少量txt标签文件中的数据以满足YOLOv5对数据集的要求；S103、将步骤S102得到的txt标签文件和对应的图片划分为训练集、验证集和测试集。
[0008]进一步地，所述步骤S2中包括：S201、将原始的YOLOv5n目标检测模型的Backbone第一层替换为conv_bn_relu_maxpool，将其下采样4倍；S202、将原始的YOLOv5n目标检测模型的Backbone除第一层外替换为ShuffleNetv2
‑
1x的stage部分，得到初步轻量化的YOLOv5n模型。
[0009]进一步地，所述步骤S3中包括：使用COCO数据集对步骤S2得到的YOLOv5n模型进行训练得到训练权重best1.pt，默认epoch为300，使用的梯度更新算法为Adam优化器，batch
‑
size为128，超参数中的学习率lr0设置为0.001。
[0010]进一步地，所述步骤S4中包括：S401、使用步骤S1得到的训练集和验证集对原始YOLOv5s进行训练，得到训练权重best2.pt；S402、修改原始YOLOv5n更新梯度的损失函数以满足知识蒸馏，其中原始YOLOv5n的损失函数为：
ꢀꢀꢀꢀꢀꢀ
(1)蒸馏损失函数为：
ꢀꢀꢀꢀ
(2)最终损失函数为：
ꢀꢀꢀꢀꢀꢀꢀꢀꢀꢀꢀꢀꢀꢀꢀꢀꢀꢀꢀꢀꢀꢀꢀꢀ
(3)其中，是真实框和预测框之间的置信度损失函数，是教师网络原始YOLOv5s的预测框与学生网络原始YOLOv5n的预测框之间的置信度损失函数，其余为分类损失函数、为目标框回归损失函数，用来平衡原始YOLOv5n的损失函数和蒸馏损失函数，这里设置为1；S403、将步骤S3得到训练权重best1.pt作为预训练权重，以及将步骤S401得到的训练权重best2.pt作为教师模型，指导步骤S2得到的YOLOv5n模型在步骤S1得到的训练集和验证集中训练，训练过程中使用步骤S402中的损失函数公式(3)进行反向传播梯度更新，训练完成得到训练权重best3.pt。
[0011]进一步地，所述步骤S5中包括：S501、将步骤S2得到的YOLOv5n本文档来自技高网...

【技术保护点】

【技术特征摘要】
1.一种基于轻量化YOLOv5的生活垃圾分类方法，其特征在于，包括以下步骤：S1、对生活垃圾数据集进行预处理，并将其划分为训练集、验证集和测试集；S2、将原始的YOLOv5n目标检测模型的Backbone替换为ShuffleNetv2的Backbone，得到初步轻量化的YOLOv5n模型；S3、使用COCO数据集对步骤S2得到的YOLOv5n模型进行训练得到训练权重best1.pt；S4、使用步骤S1选取的训练集和验证集对原始的YOLOv5s进行训练得到训练权重best2.pt，然后将best2.pt作为教师模型并使用步骤S3得到的best1.pt作为预训练权重对步骤S2得到的YOLOv5n模型进行知识蒸馏后得到训练权重best3.pt；S5、对步骤S2得到的YOLOv5n模型引入GhostNet中的Ghost模块，进一步得到轻量化的YOLOv5n模型并使用步骤S1选取的训练集和验证集对该YOLOv5n模型进行训练得到训练权重best4.pt，预训练权重为步骤S4得到的best3.pt；S6、对步骤S5得到的YOLOv5n模型引入SENet中的SELayer注意力模块得到最终轻量化的YOLOv5n模型，并使用步骤S1得到的训练集和验证集对该YOLOv5n模型进行训练得到训练权重best.pt，预训练权重为步骤S5得到的best4.pt；S7、将步骤S6的训练权重best.pt作为预训练权重，并以步骤S1选取的测试集作为步骤S6得到的最终轻量化的YOLOv5n模型的输入图像进行检测，最终检测完毕得到的图像中标明垃圾分类。2.根据权利要求1所述的基于轻量化YOLOv5的生活垃圾分类方法，其特征在于，所述步骤S1中包括：S101、使用LabelImg标注软件对生活垃圾数据集进行标签检测框调整；S102、将步骤S101调整好的数据集中的xml标签文件转换为YOLOv5可以使用的包含类别、中心点坐标、检测框高和宽的txt标签文件，其中txt标签文件中类别的编号需要连续，此外需手动修正少量txt标签文件中的数据以满足YOLOv5对数据集的要求；S103、将步骤S102得到的txt标签文件和对应的图片划分为训练集、验证集和测试集。3.根据权利要求2所述的基于轻量化YOLOv5的生活垃圾分类方法，其特征在于，所述步骤S2中包括：S201、将原始YOLOv5n的Backbone第一层替换为conv_bn_relu_maxpool，将其下采样4倍；S202、将原始YOLOv5n的Backbone除第一层外替换为ShuffleNetv2
‑
1x的stage部分，得到初步轻量化的YOLOv5n模型。4.根据权利要求3所述的基于轻量化YOLOv5的生活垃圾分类方法，其特征在于，所述步骤S3中包括：使用COCO数据集对步骤S2得到的YOLOv5n模型进行训练得到训练权重best1.pt，默认epoch为300，使用的梯度更新算法为Adam优化器，batch
‑
size为128，超参数中的学习率lr0设置为0.001。5.根据权利要求4所述的基于轻量化YOLOv5的生活垃圾分类方法，其特征在于，所述步骤S4中包括：S401、使用步骤S1得到的训练集和验证集对原始YOLOv5s进行训练，得到训练权重best2.pt；
S402、修改原始YOLOv5n更新梯度的损失...

【专利技术属性】
技术研发人员：易安林，涂成凤，姚涛，贺文伟，闫连山，蒲桂东，
申请(专利权)人：西南交通大学烟台新一代信息技术研究院，
类型：发明
国别省市：