基于改进YOLOv5的多目标垃圾分类研究
2023-08-22贺平安
孙 奥,代 琦,贺平安
(浙江理工大学,浙江 杭州 310020)
0 引言
垃圾分类是“关键小事”和“民生大事”[1]。对于垃圾的分类,我们可以通过摄像头采集垃圾图像,通过垃圾分类模型快速检测出垃圾的类别,降低人工成本,提高分类效率。因此研究垃圾分类算法有十分重要的现实意义。目前深度卷积神经网络在图像分类上的精度和速度都有很大的提升。
1989 年,LeCun 提出卷积神经网络(CNN)[2]可以有效、自动地提取图像特征,并有较快的运行速度和识别精度。在2012 年,AlexNet 获得了ImageNet 图像分类竞赛的冠军,标志着深度学习的崛起。2013 年ImageNet 图像分类竞赛的冠军由ZFNet 获得,并将top5错误率降低了近三个百分点。2014年的GoogLeNet在图像识别领域的准确度已经达到认类水平。2015年ResNet 模型的准确度超过人类水平,并且其网络层级非常深,可以提取更多的特征,标志着深度学习进入新的发展阶段。同时随着移动终端设备的不断增加,也催生了轻量级的图像识别网络,如Google 团队于2017 年发布的MobileNet,旷世团队于2017 年发布ShuffleNet,Google 团队又于2019 年发布EfficientNet。以上轻量级框架同时兼顾了精度和识别速度,方便部署在终端设备上,使得深度学习技术应用更为广泛。
我国将垃圾分为可回收垃圾、厨余垃圾、有害垃圾和其他垃圾四大类。按照此分类标准的垃圾图像识别研究还比较少,准确率低,处理速度慢,并且大部分只能识别一张图片上的单个垃圾目标。本文通过改进YOLOv5 多目标检测网络,来实现一种更轻量、快速、准确的网络模型YOLOZ。
1 网络模型及改进
1.1 YOLOv5网络
本文基于YOLOv5 网络来实现多目标的检测,YOLOv5 的Backbone 模块使用NewCSP-Darknet53 结构来进行图像的特征提取,其中ConvBNSiLU 模块用来对输入图片进行下采样。BottleNeck 模块利用了Resnet网络中的残差结构[3],解决了随着网络深度的加深而出现的网络退化问题。C3 模块将原输入分成两个分支,分别进行卷积操作,使得通道数减半,这样是为了让模型学习到更多的特征,移除计算瓶颈并降低显存的使用[4]。SPPF 是SPP(空间金字塔池化Spatial Pyramid Pooling)[5]模块的优化结构,可以加快GPU 的计算速度,并解决输入图像的大小必须固定的问题。Neck 模块使用了FPN[6]加PAN[7]的结构,FPN 就是一种自顶向下的侧边连接在不同尺度上构建出的高级语义特征图,但是这种自顶向下的FPN 网络受到单项信息流的限制,精度有待提高。PAN 是在FPN 的基础上再建立一条自下而上的通路,这样可以让高层的特征图具有更强的语义信息,有利于物体分类,底层的特征图具有更强的位置信息,有利于物体定位。Head模块的主要部分是三个Detect 检测器,利用基于网格的anchor 在不同尺度的特征图上进行目标检测,输出目标的坐标信息、置信度,以及类别信息。
1.2 BatchNormalization
Normalization 就是将所有的数据映射到同一尺度,让特征图满足均值为0,方差为1 的分布。理论上整个样本集的特征图分布都要满足这一规律,但对于大的数据集来说这是不太现实的,所以采用Batch-Normalization[8],batch 越大越接近整个数据集的分布。Batch Normalization 的计算方法如公式⑴所示,x为输入特征图,Mean 为特征图的均值,Var 标识特征图的方差,γ(默认值为1)用来调整数值分布的方差大小,β(默认值为0)用来调节数值均值的位置,这两个参数是在方向传播过程中学习得到的。
1.3 结构重参数化
本文提出一种结合VGG 和ResNet 的模块RepVGGBlock[9],在训练时使用一个类似ResNet-style的多分支模型,可以增加网络模型的表征能力防止网络出现退化现象,但是推理速度慢。而单路模型使用的3*3卷积模块推理速度快。所以我们采用结构重参数化的方法,在推理时把多分支模型修改为单路模型,如图1 所示,Conv2d 为卷积模块,BN 是BatchNormalization 模块。对于没有卷积模块的分支需要构建一个恒等的Conv2d 模块,再把所有的Conv2d 模块的卷积核大小调整为3,这样就可以进行结构重参数化得到单路结构。由公式⑵和公式⑶可知,Conv2d模块和BN模块可以融合为一个Conv2d模块。
图1 结构重参数化
1.4 BiFPN
YOLOv5 模型的Neck 结构中有些网络节点只有一条输入边,如图2左侧结构所示,删除它对我们的网络影响不大,同时可以简化网络。如果原始输入节点和输出节点在同一层,就在原始输入节点和输出节点之间添加一条额外的边,这可以在不增加太多成本的情况下融合更多的特征[10],可得到图2 右侧的结构。传统的特征融合往往只简单地把不同的特征图叠加或相加,比如Concat或Shortcut连接,而没有对不同的特征图进行区分。然而输入的各特征图分辨率也各不同,他们对融合输入特征图的贡献也是不同的,对其进行简单的相加或叠加处理并不是最佳操作,我们采用了一种类似SoftMax 的方法,把尺度归一到[0,1]之间的简单高效的加权特征融合机制。
图2 BiFPN结构
综上我们结合结构重参数化,BiFPN 模块对原始YOLO 网络进行优化,得到改进后的网络YOLOZ,如图3所示。
图3 YOLOZ网络结构
2 实验及结果分析
2.1 实验环境
本文实验在Windows10系统下,利用Python语言,Pytorch 深度学习框架。机器配置为CPUIntelCorei5-11600k@3.9GHz,内存为32GB,显卡为NvidiaGeForceRTX3050,显存为8GB。
2.2 实验数据
数据来源为华为云AI 大赛提供的垃圾图片数据以及网络收集的垃圾图片,并手动划分四大类,分别为:可回收垃圾、有害垃圾、易腐垃圾、其他垃圾。每大类又细分为三小类。其中可回收垃圾细分为纸箱、塑料瓶、泡沫盒子,有害垃圾细分为电池、灯泡、软膏,易腐垃圾细分为剩饭剩菜、蛋糕、蛋壳,其他垃圾细分为牙刷、口罩、烟蒂。共12 个类别3932 张图片,8448个标注信息。
2.3 实验结果
2.3.1 Mosaic数据增强
为了增强数据集的图片和标签数量,可采用Mosaic数据增强的方式,Mosaic的方法是:使用4张图片,对其进行随机裁剪和缩放旋转等操作,最后合成1张图像,原图的标签信息也映射到新生成的图像中,丰富了检测物体的背景,增加了小尺度目标的数量、数据集图片的数量以及标签信息。但是Mosaic 数据增强是一个耗时的方法,所以我们可以设置一个阈值和一个随机数,当生成的随机数小于该阈值时才采用Mosaic 数据增强,来根据不同的场景权衡训练速度和模型精度的平衡。表1 所示是设置不同阈值下的Mosaic数据增强对模型精度的影响。
表1 Mosaic数据增强结果
由表1可知,本实验数据较少,如果不采用Mosaic数据增强的方式,模型的精度很低,采用Mosaic 数据增强时,模型的精度相比不使用Mosaic 数据增强大幅提高,且随着阈值的增大,精度不断提高,但是模型每轮训练耗时也再增加。
2.3.2 结构重参数化
本论文模型中使用了结构重参数化技术,可以让训练后的模型在推理时保持精度不变的前提下提高推理速度。表2比较了同一模型在使用该技术前后的推理速度和推理精度。其中图片的输入尺寸全部被调整为640*640 大小,每张图片的推理时间是推理测试集所需时间除以测试集图片数的结果。
表2 结构重参数化
由表2可知,开启结构重参数化后,网络的推理速度提升了8%,模型的计算量下降了3.2%,并且网络的精度没有下降。
2.3.3 BiFPN
本论文模型中引进了BiFPN 网络,优化了原模型中的FPN 和PAN 结构,可以在不增加太多成本的情况下融合更多特征。通过实验对比使用BiFPN 优化前后的网络,结果如表3所示。
表3 BiFPN结构优化网络
由表3 可知,当使用了BiFPN 时网络的mAP@.5提高了1.1 个百分点,mAP@.5:.95 提高了2 个百分点,引入BiFPN 结构相比于原始的FPN+PAN 可以提高网络的精度。
2.3.4 不同多目标检测模型对比
基于YOLOv5 并结合本文所述的三种对网络的改进得到YOLOZ 网络,通过和前几代YOLO 版本以及其他目标检测框架的比较,得出结果如表4 所示。可见相比较于YOLOv3 和YOLOv4,模型的大小显著降低,并且精度也有所提高。相较于YOLOv5 系列,本模型体积比YOLOv5s略大,但是在精度上有很大提高,并且比YOLOv5m 体积更小,精度更高。相比较于其他目标检测模型如EfficientDet,在检测速度和模型大小方面都占据优势。原YOLOv5 模型与YOLOZ模型在训练时的mAP 曲线如图4 所示,可见YOLOZ模型在收敛速度与mAP值上均优于原YOLOv5模型。
表4 不同目标检测网络的对比
图4 模型修改前后训练时mAP对比
3 结论
本文针对多目标垃圾分类问题改进了YOLOv5网络的Backbone 结构、Neck 结构,并对数据集采用了Mosaic 的增强方法。得到了YOLOZ 模型,该模型在精度和体积方面都优于改进前的模型。但是在数据集的数量与标注质量上,本数据集图片质量较低,数量较少,且分类类别较多,导致模型的mAP 值不高。接下来,还需要继续对YOLOZ 模型进行优化,结合ShuffleNet、EfficientNet 等轻量级网络,进一步减小模型体积,并将其移植到嵌入式设备中。