CN115439684B - 基于轻量化YOLOv5的生活垃圾分类方法及APP - Google Patents
基于轻量化YOLOv5的生活垃圾分类方法及APP Download PDFInfo
- Publication number
- CN115439684B CN115439684B CN202211028322.8A CN202211028322A CN115439684B CN 115439684 B CN115439684 B CN 115439684B CN 202211028322 A CN202211028322 A CN 202211028322A CN 115439684 B CN115439684 B CN 115439684B
- Authority
- CN
- China
- Prior art keywords
- training
- yolov5n
- model
- lightweight
- household garbage
- Prior art date
- Legal status (The legal status is an assumption and is not a legal conclusion. Google has not performed a legal analysis and makes no representation as to the accuracy of the status listed.)
- Active
Links
- 239000010813 municipal solid waste Substances 0.000 title claims abstract description 77
- 238000000034 method Methods 0.000 title claims abstract description 35
- 238000001514 detection method Methods 0.000 claims abstract description 50
- 238000013140 knowledge distillation Methods 0.000 claims abstract description 13
- 238000012549 training Methods 0.000 claims description 125
- 230000006870 function Effects 0.000 claims description 34
- 238000012795 verification Methods 0.000 claims description 27
- 238000004422 calculation algorithm Methods 0.000 claims description 15
- 101100004280 Caenorhabditis elegans best-2 gene Proteins 0.000 claims description 12
- 101100493820 Caenorhabditis elegans best-1 gene Proteins 0.000 claims description 11
- 238000012360 testing method Methods 0.000 claims description 9
- 238000004821 distillation Methods 0.000 claims description 6
- ORILYTVJVMAKLC-UHFFFAOYSA-N Adamantane Natural products C1C(C2)CC3CC1CC2C3 ORILYTVJVMAKLC-UHFFFAOYSA-N 0.000 claims description 3
- 238000001914 filtration Methods 0.000 claims description 3
- 238000007781 pre-processing Methods 0.000 claims description 3
- 230000001629 suppression Effects 0.000 claims description 3
- 238000013461 design Methods 0.000 abstract description 5
- 238000004364 calculation method Methods 0.000 abstract description 4
- 238000010586 diagram Methods 0.000 description 6
- 230000006835 compression Effects 0.000 description 4
- 238000007906 compression Methods 0.000 description 4
- 210000000988 bone and bone Anatomy 0.000 description 3
- 238000013145 classification model Methods 0.000 description 2
- 238000013527 convolutional neural network Methods 0.000 description 2
- 238000005516 engineering process Methods 0.000 description 2
- 238000000605 extraction Methods 0.000 description 2
- 238000013139 quantization Methods 0.000 description 2
- 238000011160 research Methods 0.000 description 2
- 238000004088 simulation Methods 0.000 description 2
- 239000002699 waste material Substances 0.000 description 2
- 238000013473 artificial intelligence Methods 0.000 description 1
- 230000009286 beneficial effect Effects 0.000 description 1
- 238000013135 deep learning Methods 0.000 description 1
- 230000007547 defect Effects 0.000 description 1
- 230000001934 delay Effects 0.000 description 1
- 230000000694 effects Effects 0.000 description 1
- 238000002474 experimental method Methods 0.000 description 1
- 239000010806 kitchen waste Substances 0.000 description 1
- 238000003062 neural network model Methods 0.000 description 1
- 238000012805 post-processing Methods 0.000 description 1
- 238000013138 pruning Methods 0.000 description 1
- 230000000007 visual effect Effects 0.000 description 1
Classifications
-
- G—PHYSICS
- G06—COMPUTING; CALCULATING OR COUNTING
- G06V—IMAGE OR VIDEO RECOGNITION OR UNDERSTANDING
- G06V10/00—Arrangements for image or video recognition or understanding
- G06V10/70—Arrangements for image or video recognition or understanding using pattern recognition or machine learning
- G06V10/764—Arrangements for image or video recognition or understanding using pattern recognition or machine learning using classification, e.g. of video objects
-
- G—PHYSICS
- G06—COMPUTING; CALCULATING OR COUNTING
- G06N—COMPUTING ARRANGEMENTS BASED ON SPECIFIC COMPUTATIONAL MODELS
- G06N3/00—Computing arrangements based on biological models
- G06N3/02—Neural networks
- G06N3/08—Learning methods
- G06N3/084—Backpropagation, e.g. using gradient descent
-
- G—PHYSICS
- G06—COMPUTING; CALCULATING OR COUNTING
- G06V—IMAGE OR VIDEO RECOGNITION OR UNDERSTANDING
- G06V10/00—Arrangements for image or video recognition or understanding
- G06V10/70—Arrangements for image or video recognition or understanding using pattern recognition or machine learning
- G06V10/77—Processing image or video features in feature spaces; using data integration or data reduction, e.g. principal component analysis [PCA] or independent component analysis [ICA] or self-organising maps [SOM]; Blind source separation
- G06V10/774—Generating sets of training patterns; Bootstrap methods, e.g. bagging or boosting
-
- G—PHYSICS
- G06—COMPUTING; CALCULATING OR COUNTING
- G06V—IMAGE OR VIDEO RECOGNITION OR UNDERSTANDING
- G06V10/00—Arrangements for image or video recognition or understanding
- G06V10/70—Arrangements for image or video recognition or understanding using pattern recognition or machine learning
- G06V10/82—Arrangements for image or video recognition or understanding using pattern recognition or machine learning using neural networks
-
- G—PHYSICS
- G06—COMPUTING; CALCULATING OR COUNTING
- G06V—IMAGE OR VIDEO RECOGNITION OR UNDERSTANDING
- G06V2201/00—Indexing scheme relating to image or video recognition or understanding
- G06V2201/07—Target detection
Landscapes
- Engineering & Computer Science (AREA)
- Theoretical Computer Science (AREA)
- Evolutionary Computation (AREA)
- Physics & Mathematics (AREA)
- Artificial Intelligence (AREA)
- General Physics & Mathematics (AREA)
- Computing Systems (AREA)
- Health & Medical Sciences (AREA)
- General Health & Medical Sciences (AREA)
- Software Systems (AREA)
- Databases & Information Systems (AREA)
- Medical Informatics (AREA)
- Computer Vision & Pattern Recognition (AREA)
- Multimedia (AREA)
- Life Sciences & Earth Sciences (AREA)
- Biomedical Technology (AREA)
- Biophysics (AREA)
- Computational Linguistics (AREA)
- Data Mining & Analysis (AREA)
- Molecular Biology (AREA)
- General Engineering & Computer Science (AREA)
- Mathematical Physics (AREA)
- Management, Administration, Business Operations System, And Electronic Commerce (AREA)
Abstract
本发明提供了一种基于轻量化YOLOv5的生活垃圾分类方法及APP，面向移动端设备手机设计部署，以原始的目标检测网络YOLOv5n为基线网络，采用引入ShuffleNetv2中的shuffle_block、GhostNet的Ghost模块、引入注意力机制SE以及知识蒸馏等技术手段，结合网络结构设计和知识蒸馏对基线网络实现轻量化，轻量化后的YOLOv5n模型相比原始的YOLOv5n模型，在很大程度上减少了模型的参数量和计算量以及提升了检测精度，将轻量化后的YOLOv5模型部署至移动端的生活垃圾分类APP中，可以实现对多种生活垃圾智能分类。
Description
技术领域
本发明属于计算机视觉技术领域，具体涉及一种基于轻量化YOLOv5的生活垃圾分类方法及APP。
背景技术
近年来，计算机视觉一直是科学界与工业界的研究热门点，作为计算机视觉中的基本视觉识别问题，目标检测在过去的几十年中得到了广泛地研究。目标检测是对图像或视频中的目标进行分类和定位的任务。目前，目标检测方法主要分为两类，分别是两阶段目标检测方法和一阶段目标检测方法。两阶段目标检测方法，例如FasterR-CNN和CascadeR-CNN和DetectoRS等。一阶段目标检测方法典型代表是YOLO系列以及SSD。两类目标检测方法各有优势：两阶段算法相对而言精度更高，尤其体现在定位上，而一阶段算法的速度普遍更快，更容易满足实际应用场景中的实时性需求。
目前人工智能技术应用的另一个重要趋势是在端侧平台上部署高性能的神经网络模型并能在真实场景中实时运行，如移动端/嵌入式端设备。然而现有的这些端侧平台的特点是内存资源少，处理器性能不高，功耗受限，这使得目前精度最高的模型由于对内存和计算资源的超额要求，使得其根本无法在现有的端侧平台上面部署且达到实时性的要求。为了使模型落地应用，需要设计轻量型网络结构或者通过知识蒸馏（DeepCompression）、剪枝(NetWorkPruning)、量化(ModelQuantization)等技术减少模型参数、计算量。大多数的模型压缩方法是在分类模型上进行的。常见的分类模型通常输入分辨率小、模型结构简单，而目标检测模型通常有着更大的输入分辨率，同时网络结构也更加复杂。目标检测模型可能还包括多个输入，以及一些后处理过程，这些都对模型压缩算法的设计提出了更高的要求。如果要在移动端部署目标检测网络，通常需要更浅的模型深度，更少的卷积核，但是不够深的网络模型往往在特征提取上不充分。因此一个平衡精度，检测速度和模型复杂度的目标检测网络显得至关重要。
目前生活垃圾分类主要面临以下问题：一是垃圾分类还是以人工分拣为主，因为人主观判断的不同造成分类错误从而导致无法从源头实现垃圾分类；二是智能分类垃圾桶尚未普及，而且现有智能垃圾分类系统一般采用的模型算法无法同时对多种垃圾进行分类，效率低下。此外，目前基于深度学习的生活垃圾分类方法针对的垃圾类别较少。如果提出一种方法能够快速便捷的区分日常生活垃圾属于可回收物、有害垃圾、厨余垃圾还是其他垃圾，就能够很大提高生活垃圾分类效率以及积极性。
发明内容
本发明的目的在于针对上述现有技术中的不足，提供一种基于轻量化YOLOv5的生活垃圾分类方法及APP，将轻量化目标检测网络在手机部署以实现日常生活垃圾的分类。
本发明采用以下技术方案实现：
一种基于轻量化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模型的输入图像进行检测，最终检测完毕得到的图像中标明垃圾分类。
进一步地，所述步骤S1中包括：
S101、使用LabelImg标注软件对生活垃圾数据集进行标签检测框调整；
S102、将步骤S101调整好的数据集中的xml标签文件转换为YOLOv5可以使用的包含类别、中心点坐标、检测框高和宽的txt标签文件，其中txt标签文件中类别的编号需要连续，此外需手动修正少量txt标签文件中的数据以满足YOLOv5对数据集的要求；
S103、将步骤S102得到的txt标签文件和对应的图片划分为训练集、验证集和测试集。
进一步地，所述步骤S2中包括：
S201、将原始的YOLOv5n目标检测模型的Backbone第一层替换为conv_bn_relu_maxpool，将其下采样4倍；
S202、将原始的YOLOv5n目标检测模型的Backbone除第一层外替换为ShuffleNetv2-1x的stage部分，得到初步轻量化的YOLOv5n模型。
进一步地，所述步骤S3中包括：
使用COCO数据集对步骤S2得到的YOLOv5n模型进行训练得到训练权重best1.pt，默认epoch为300，使用的梯度更新算法为Adam优化器，batch-size为128，超参数中的学习率lr0设置为0.001。
进一步地，所述步骤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。
进一步地，所述步骤S5中包括：
S501、将步骤S2得到的YOLOv5n模型的Neck中的Conv和C3替换为GhostNet中的Ghost模块以及由Ghost模块组成的C3Ghost模块；
S502、在步骤S501的基础上将下采样层设置为5×5的Ghost卷积核；
S503、使用步骤S1得到的训练集和验证集对步骤S502得到的进一步轻量化的YOLOv5n模型进行训练得到训练权重best4.pt，训练过程使用的预训练权重为步骤S403得到的best3.pt。
进一步地，所述步骤S6中包括：
S601、在步骤S502得到的YOLOv5n模型基础上，引入SENet的SELayer模块并将其添加至该YOLOv5n模型的Backbone最后一层，得到最终轻量化的YOLOv5n模型；
S602、使用步骤S1得到的训练集和验证集对步骤S601得到的最终轻量化YOLOv5n模型进行训练得到训练权重best.pt，预训练权重为步骤S503得到的best4.pt。
进一步地，所述步骤S7中包括：
在对输入图像进行检测时，进行前向推理使用置信度阈值以及非极大抑制算法进行检测框的过滤，最终检测完毕得到的图像中的垃圾分类包含有检测框、类别名称（例如papercup、can等）、置信度以及该图像中的垃圾是可回收物、厨余垃圾、有害垃圾还是其他垃圾。
进一步地，本发明还提供一种基于轻量化YOLOv5的生活垃圾分类APP，包括下列步骤：
A，将上述所述方法中的步骤S6得到的训练权重best.pt转换为best.onnx，并进一步使用onnx-simpifier工具将其简化为best-sim.onnx；
B，使用ncnn源码编译配置中的onnx2ncnn.exe将步骤A的best-sim.onnx转换为YOLOv5n网络结构文件xxx.param和YOLOv5n网络权重文件xxx.bin；
C，将步骤B的网络结构文件xxx.param和网络权重文件xxx.bin部署至ncnn针对Android写的demo项目中，并修改相应的代码，完成APP部署。
与现有技术相比，本发明具有以下有益效果：
本发明是以原始的目标检测网络YOLOv5n为基线网络，采用引入ShuffleNetv2中的shuffle_block、GhostNet的Ghost模块、引入注意力机制SE以及知识蒸馏等技术手段，结合网络结构设计和知识蒸馏对基线网络实现轻量化，轻量化后的YOLOv5n模型相比原始的YOLOv5n模型，在很大程度上减少了模型的参数量和计算量以及提升了检测精度，因此可以将本发明的应用技术部署至端侧平台使用。本发明还提供一种基于轻量化YOLOv5的生活垃圾分类APP，将本发明的生活垃圾分类APP部署至手机端，可以实现对多种生活垃圾智能分类。居民使用智能生活垃圾分类APP不仅可以提高对生活垃圾进行分类的准确性，而且降低了对生活垃圾进行分类的难度，进一步提高群众对生活垃圾分类的兴趣，有利于国家推广生活垃圾分类政策，从源头实现垃圾正确分类，更好的保护地球环境和节省地球资源。
附图说明
图1为shuffle_block图。
图2为Ghost模块图。
图3为SELayer模块图。
图4为本发明的改进后的轻量化YOLOv5n模型框架图。
图5为应用本发明生活垃圾分类APP对垃圾进行分类的结果示意图。
下面结合附图对本发明的技术方案进行清楚、完整地描述，显然，所描述的实施例是本发明一部分实施例，而不是全部的实施例。基于本发明中的实施例，本领域普通技术人员在没有做出创造性劳动前提下所获得的所有其他实施例，都属于本发明保护的范围。
具体实施方式
本发明提供了一种基于轻量化YOLOv5的生活垃圾分类方法及APP，面向移动端设备手机设计部署，结合人工设计轻量型网络结构和模型压缩算法中的知识蒸馏，可以实现多种生活垃圾检测分类，在满足可接受推理延时前提下，降低模型的参数量、计算量，提升了检测精度。
实施例1
一种基于轻量化YOLOv5的生活垃圾分类方法，包括以下步骤：
S1、对生活垃圾数据集进行预处理，并将其划分为训练集、验证集和测试集，具体包括：
S101、采用包含样本数量为18170的30类生活垃圾数据集HGI_30_detection，使用LabelImg标注软件对上述生活垃圾数据集进行标签检测框调整；
S102、将步骤S101调整好的数据集中的xml标签文件转换为YOLOv5可以使用的包含类别、中心点坐标、检测框高和宽的txt标签文件，其中txt标签文件中类别的编号连续由0-29，同时手动修正少量txt标签文件的数据以满足YOLOv5对数据集的要求；
S103、将步骤S102得到的txt标签文件和对应的图片选取170张图片作为测试集，剩余标签和图片以8:2比例划分为样本数为14400的训练集、样本数为3600的验证集。
S2、采用一阶段目标检测算法代表作YOLOv5为研究基础模型，将原始的YOLOv5n目标检测模型的Backbone替换为ShuffleNetv2的Backbone，得到初步轻量化的YOLOv5n模型，具体包括：
S201、将原始YOLOv5n的Backbone第一层替换为conv_bn_relu_maxpool，将其下采样4倍，通道增加至32；
S202、将原始YOLOv5n的Backbone除第一层外替换为ShuffleNetv2-1x的stage部分，得到初步轻量化的YOLOv5n模型；
需要知道步骤S202所提到的ShuffleNetv2-1x的stage部分是如图1所示的shuffle_block模块，该模块不仅包含步长为2的下采样单元，还包括步长为1的提取特征单元。
S3、使用COCO数据集对步骤S2得到的YOLOv5n模型进行训练得到训练权重best1.pt，默认epoch为300，使用的梯度更新算法为Adam优化器，batch-size为128，超参数中的学习率lr0设置为0.001。
S4、使用步骤S1选取的训练集和验证集对原始的YOLOv5s进行训练得到训练权重best2.pt，然后将best2.pt作为教师模型并使用步骤S3得到的best1.pt作为预训练权重对步骤S2得到的YOLOv5n模型进行知识蒸馏后得到训练权重best3.pt，具体包括：
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。
S5、对步骤S2得到的YOLOv5n模型引入GhostNet中的Ghost模块，进一步得到轻量化的YOLOv5n模型，并使用步骤S1选取的训练集和验证集对该YOLOv5n模型进行训练得到训练权重best4.pt，预训练权重为步骤S4得到的best3.pt，具体包括；
S501、将步骤S2得到的YOLOv5n模型的Neck中的Conv和C3替换为GhostNet中的Ghost模块以及由Ghost模块组成的C3Ghost模块，参见图2所示的Ghost模块；
S502、在步骤S501的基础上将下采样层设置为5×5的Ghost卷积核；
S503、使用步骤S1得到的训练集和验证集对步骤S502得到的进一步轻量化的YOLOv5n模型进行训练得到训练权重best4.pt，训练过程使用的预训练权重为步骤S403得到的best3.pt。
S6、对步骤S5得到的YOLOv5n模型引入SENet中的SELayer注意力模块（参见图3所示的SELayer模块）得到最终轻量化的YOLOv5n模型，并使用步骤S1得到的训练集和验证集对该YOLOv5n模型进行训练得到训练权重best.pt，预训练权重为步骤S5得到的best4.pt，具体包括：
S601、在步骤S502得到的YOLOv5n模型基础上，引入SENet的SELayer模块并将其添加至该YOLOv5n模型的Backbone最后一层，得到最终轻量化的YOLOv5n模型（参见图4所示）；
S602、使用步骤S1得到的训练集和验证集对步骤S601得到的最终轻量化YOLOv5n模型进行训练得到训练权重best.pt，预训练权重为步骤S503得到的best4.pt。
S7、将步骤S6的训练权重best.pt作为预训练权重，并以步骤S1选取的测试集作为步骤S6得到的最终轻量化的YOLOv5n模型的输入图像进行检测，检测时进行前向推理使用置信度阈值以及非极大抑制算法（NMS）进行检测框的过滤，最终检测完毕得到的图像中标明垃圾分类，垃圾分类包含有检测框、类别名称（例如papercup、can等）、置信度以及该图像中的垃圾是可回收物、厨余垃圾、有害垃圾还是其他垃圾。
实施例2
一种基于轻量化YOLOv5的生活垃圾分类APP，包括以下步骤；
A，将上述所述方法中的步骤S6得到的训练权重best.pt转换为best.onnx，并进一步使用onnx-simpifier工具将其简化为best-sim.onnx；
B，使用ncnn源码编译配置中的onnx2ncnn.exe将步骤A的best-sim.onnx转换为YOLOv5n网络结构文件xxx.param和YOLOv5n网络权重文件xxx.bin；
C，将步骤B的网络结构文件xxx.param和网络权重文件xxx.bin部署至ncnn针对Android手机写的demo项目中，并修改相应的代码，完成APP部署；
D，选中手机相册中的一张图片或者结合opencv实现打开摄像头进行场景拍照，输出对图片识别后的目标检测框、类别、置信度以及所属可回收物、厨余垃圾、有害垃圾还是其他垃圾中的某一类，参见图5所示。
本发明的效果可以通过以下实验进一步说明。
1、实验环境：
训练主机仿真环境为：WindowsServer2012Intel(R)Xeon(R) CPUE52360，内存为128G，GPU为NVIDIATITANRTX24G；
采用华为Honor30，安卓端推理环境：处理器为HUAWEIKirin985，运行内存为8G，GPU为Turbo 。
2、实验内容：
以原始YOLOv5n和原始YOLOv5s为基线模型，将原始YOLOv5n的Backbone替换为更轻量的Shufflenetv2；将原始YOLOv5s作为教师模型，通过Response-Based KnowledgeDistillation算法提升替换骨干后的YOLOv5n的检测精度；此外为了将其进一步轻量化，引入了Ghost模块和C3Ghost模块；最后引入注意力机制SELayer模块提升轻量化后的YOLOv5n的检测精度。记录每次引入新模块或算法后YOLOv5n的参数量、计算量以及mAP0.5和mAP0.5:0.95，最后使用ncnn推理框架将原始YOLOv5n、YOLOv5s模型以及最终轻量化的YOLOv5n模型部署至安卓端Honor30的生活垃圾分类APP中，验证模型的推理延时。
3、仿真实验结果：
表1为引入各个模块和算法的YOLOv5n与原始YOLOv5n、原始YOLOv5s的性能对比，下表所有模型的浮点精度均为32。
表2为原始YOLOv5n、轻量化的YOLOv5n、原始YOLOv5s在安卓端Honor30上的推理延时对比。
可以看到，经过轻量化的YOLOv5n模型的参数量、计算量以及模型文件大小都有显著减少，相比原始YOLOv5n模型，参数量下降了22.9%，计算量下降了23.3%，mAP0.5由0.879增加至0.895，且推理延时相对原始YOLOv5n模型而言只增加了约6ms，相比原始YOLOv5s模型，参数量下降了80.5%，计算量下降了80.5%，mAP0.5由0.89增加至0.895，完全可以满足将模型部署至安卓端并进行生活垃圾分类的需求，而且将部署之前的浮点精度为32转为16可以进一步减少推理延时。
Claims (9)
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更新梯度的损失函数以满足知识蒸馏，其中原始YOLOv5n的损失函数为： (1)
蒸馏损失函数为：
(2)
最终损失函数为：
(3)
其中，是真实框和预测框之间的置信度损失函数，/>是教师网络原始YOLOv5s的预测框与学生网络原始YOLOv5n的预测框之间的置信度损失函数，其余为分类损失函数、为目标框回归损失函数，用来平衡原始YOLOv5n的损失函数和蒸馏损失函数，这里设置为1；
S403、将步骤S3得到的训练权重best1.pt作为预训练权重，以及将步骤S401得到的训练权重best2.pt作为教师模型，指导步骤S2得到的YOLOv5n模型在步骤S1得到的训练集和验证集中训练，训练过程中使用步骤S402中的损失函数公式(3)进行反向传播梯度更新，训练完成得到训练权重best3.pt。
6.根据权利要求5所述的基于轻量化YOLOv5的生活垃圾分类方法，其特征在于，所述步骤S5中包括：
S501、将步骤S2得到的YOLOv5n模型的Neck中的Conv和C3替换为GhostNet中的Ghost模块以及由Ghost模块组成的C3Ghost模块；
S502、在步骤S501的基础上将下采样层设置为5×5的Ghost卷积核；
S503、使用步骤S1得到的训练集和验证集对步骤S502得到的进一步轻量化的YOLOv5n模型进行训练得到训练权重best4.pt，训练过程使用的预训练权重为步骤S403得到的best3.pt。
7.根据权利要求6所述的基于轻量化YOLOv5的生活垃圾分类方法，其特征在于，所述步骤S6中包括：
S601、在步骤S502得到的YOLOv5n模型基础上，引入SENet的SELayer模块并将其添加至该YOLOv5n模型的Backbone最后一层，得到最终轻量化的YOLOv5n模型；
S602、使用步骤S1得到的训练集和验证集对步骤S601得到的最终轻量化YOLOv5n模型进行训练得到训练权重best.pt，预训练权重为步骤S503得到的best4.pt。
8.根据权利要求1所述的基于轻量化YOLOv5的生活垃圾分类方法，其特征在于，所述步骤S7中包括：
在对输入图像进行检测时，进行前向推理使用置信度阈值以及非极大抑制算法进行检测框的过滤，最终检测完毕得到的图像中的垃圾分类包含有检测框、类别名称、置信度以及该图像中的垃圾是可回收物、厨余垃圾、有害垃圾还是其他垃圾。
9.一种基于轻量化YOLOv5的生活垃圾分类系统，其特征在于，包括下列步骤：
A，将上述权利要求1所述方法中的步骤S6得到的训练权重best.pt转换为best.onnx，并进一步使用onnx-simpifier工具将其简化为best-sim.onnx；
B，使用ncnn源码编译配置中的onnx2ncnn.exe将步骤A的best-sim.onnx转换为YOLOv5n网络结构文件xxx.param和YOLOv5n网络权重文件xxx.bin；
C，将步骤B的网络结构文件xxx.param和网络权重文件xxx.bin部署至ncnn针对Android写的demo项目中，并修改相应的代码，完成系统部署。
Priority Applications (1)
| Application Number | Priority Date | Filing Date | Title |
|---|---|---|---|
| CN202211028322.8A CN115439684B (zh) | 2022-08-25 | 2022-08-25 | 基于轻量化YOLOv5的生活垃圾分类方法及APP |
Applications Claiming Priority (1)
| Application Number | Priority Date | Filing Date | Title |
|---|---|---|---|
| CN202211028322.8A CN115439684B (zh) | 2022-08-25 | 2022-08-25 | 基于轻量化YOLOv5的生活垃圾分类方法及APP |
Publications (2)
| Publication Number | Publication Date |
|---|---|
| CN115439684A CN115439684A (zh) | 2022-12-06 |
| CN115439684B true CN115439684B (zh) | 2024-02-02 |
Family
ID=84245207
Family Applications (1)
| Application Number | Title | Priority Date | Filing Date |
|---|---|---|---|
| CN202211028322.8A Active CN115439684B (zh) | 2022-08-25 | 2022-08-25 | 基于轻量化YOLOv5的生活垃圾分类方法及APP |
Country Status (1)
| Country | Link |
|---|---|
| CN (1) | CN115439684B (zh) |
Citations (8)
| Publication number | Priority date | Publication date | Assignee | Title |
|---|---|---|---|---|
| CN111213160A (zh) * | 2017-11-14 | 2020-05-29 | 谷歌有限责任公司 | 高效卷积神经网络 |
| CN113011365A (zh) * | 2021-03-31 | 2021-06-22 | 中国科学院光电技术研究所 | 一种结合轻量化网络的目标检测方法 |
| CN113469073A (zh) * | 2021-07-06 | 2021-10-01 | 西安电子科技大学 | 一种基于轻量级深度学习的sar图像舰船检测方法及系统 |
| CN113469264A (zh) * | 2021-07-14 | 2021-10-01 | 南通通琦自动化科技有限公司 | 一种垃圾自动分类模型的构建方法、垃圾分拣方法和系统 |
| CN114170512A (zh) * | 2021-12-08 | 2022-03-11 | 西安电子科技大学 | 基于网络剪枝和参数量化结合的遥感sar目标检测方法 |
| CN114332666A (zh) * | 2022-03-11 | 2022-04-12 | 齐鲁工业大学 | 一种基于轻量化神经网络模型的图像目标检测方法及系统 |
| CN114862768A (zh) * | 2022-04-14 | 2022-08-05 | 长沙理工大学 | 基于改进YOLOv5-LITE轻量级的配电组件缺陷识别方法 |
| CN114926907A (zh) * | 2022-06-06 | 2022-08-19 | 金陵科技学院 | 一种基于手势识别的无接触式电梯控制方法 |
-
2022
- 2022-08-25 CN CN202211028322.8A patent/CN115439684B/zh active Active
Patent Citations (8)
| Publication number | Priority date | Publication date | Assignee | Title |
|---|---|---|---|---|
| CN111213160A (zh) * | 2017-11-14 | 2020-05-29 | 谷歌有限责任公司 | 高效卷积神经网络 |
| CN113011365A (zh) * | 2021-03-31 | 2021-06-22 | 中国科学院光电技术研究所 | 一种结合轻量化网络的目标检测方法 |
| CN113469073A (zh) * | 2021-07-06 | 2021-10-01 | 西安电子科技大学 | 一种基于轻量级深度学习的sar图像舰船检测方法及系统 |
| CN113469264A (zh) * | 2021-07-14 | 2021-10-01 | 南通通琦自动化科技有限公司 | 一种垃圾自动分类模型的构建方法、垃圾分拣方法和系统 |
| CN114170512A (zh) * | 2021-12-08 | 2022-03-11 | 西安电子科技大学 | 基于网络剪枝和参数量化结合的遥感sar目标检测方法 |
| CN114332666A (zh) * | 2022-03-11 | 2022-04-12 | 齐鲁工业大学 | 一种基于轻量化神经网络模型的图像目标检测方法及系统 |
| CN114862768A (zh) * | 2022-04-14 | 2022-08-05 | 长沙理工大学 | 基于改进YOLOv5-LITE轻量级的配电组件缺陷识别方法 |
| CN114926907A (zh) * | 2022-06-06 | 2022-08-19 | 金陵科技学院 | 一种基于手势识别的无接触式电梯控制方法 |
Non-Patent Citations (1)
| Title |
|---|
| Garbage Classification Detection Model Based on YOLOv4 with Lightweight Neural Network Feature;Xiao-Feng Wang等;International Conference on Intelligent Computing;第435-447页 * |
Also Published As
| Publication number | Publication date |
|---|---|
| CN115439684A (zh) | 2022-12-06 |
Similar Documents
| Publication | Publication Date | Title |
|---|---|---|
| US20210042580A1 (en) | Model training method and apparatus for image recognition, network device, and storage medium | |
| CN109886359B (zh) | 基于卷积神经网络的小目标检测方法及检测系统 | |
| CN111950723B (zh) | 神经网络模型训练方法、图像处理方法、装置及终端设备 | |
| US11138903B2 (en) | Method, apparatus, device and system for sign language translation | |
| CN110223292B (zh) | 图像评估方法、装置及计算机可读存储介质 | |
| CN112101437A (zh) | 基于图像检测的细粒度分类模型处理方法、及其相关设备 | |
| CN109685713B (zh) | 化妆模拟控制方法、装置、计算机设备及存储介质 | |
| CN110908518B (zh) | 基于虚拟现实的数学教学系统 | |
| CN110413551B (zh) | 信息处理装置、方法及设备 | |
| CN113569737A (zh) | 基于自主学习网络模型笔记本屏幕缺陷检测方法及介质 | |
| CN110533074B (zh) | 一种基于双深度神经网络的图片类别自动标注方法及系统 | |
| CN115439684B (zh) | 基于轻量化YOLOv5的生活垃圾分类方法及APP | |
| CN112560668B (zh) | 一种基于场景先验知识的人体行为识别方法 | |
| CN113743203A (zh) | 基于深度迁移学习网络的笔记本屏幕缺陷检测方法及设备 | |
| CN109298783A (zh) | 基于表情识别的标注监控方法、装置及电子设备 | |
| CN116797864B (zh) | 基于智能镜子的辅助化妆方法、装置、设备及存储介质 | |
| CN117456246A (zh) | 基于ViT的类内、类间相似度的细粒度图像分类方法及系统 | |
| CN109636711B (zh) | 漫画册生成方法、装置及计算机可读存储介质 | |
| CN111652242A (zh) | 图像处理方法、装置、电子设备及存储介质 | |
| CN116128044A (zh) | 一种模型剪枝方法、图像处理方法及相关装置 | |
| CN110231960A (zh) | 一种目标屏幕确定方法、装置及存储介质 | |
| CN113469254A (zh) | 基于目标检测模型的目标检测方法及系统 | |
| CN111539248B (zh) | 一种红外线人脸检测方法、装置及其电子设备 | |
| CN114120096A (zh) | 一种基于交互式ai的红绿灯标注方法及系统 | |
| TWI793951B (zh) | 模型訓練方法與模型訓練系統 |
Legal Events
| Date | Code | Title | Description |
|---|---|---|---|
| PB01 | Publication | ||
| PB01 | Publication | ||
| SE01 | Entry into force of request for substantive examination | ||
| SE01 | Entry into force of request for substantive examination | ||
| GR01 | Patent grant | ||
| GR01 | Patent grant |