基于改进YOLOv4的垃圾识别研究

胡亚兰 柯敏毅

（湖北工业大学计算机学院 湖北省武汉市 430068）

1 引言

随着社会生活水平的不断提高，人们购入了丰富的物资，越来越多的垃圾也随即产出，而垃圾种类繁多，每个种类的处理方式也不一样。例如：有害垃圾采取热处理、化学分解和无污染处理的形式；采用生化处理、沼泽生产、堆肥等方式对湿垃圾进行资源化利用或无害化处理； 干垃圾通过焚烧方式无害化处理； 可回收垃圾运至相关企业进行回收处理[1]。由此，衍生出了两个社会问题：

（1）由于人们从小没有垃圾分类的意识，容易将垃圾混在一起丢进垃圾箱，这给垃圾分拣场的工作人员带来极大困扰。垃圾分炼厂需要消耗极大的人力、时间去将不同种类的垃圾分到对应类别的垃圾框中。由于垃圾分类处理效率不高，导致一部分垃圾未被处理，造成了环境的污染和资源的浪费。

（2）垃圾分类问题逐渐引起社会的关注，2019年7月1日，上海实施垃圾分类政策，将垃圾分为可回收垃圾、厨余垃圾、有害垃圾和其它垃圾四大类。但是，由于人们垃圾分类意识比较薄弱，容易将垃圾分错类别，还是会出现将垃圾混在一起的情况。这样仍然会使一部分垃圾处于未处理状态，给垃圾处理场的垃圾分类、垃圾处理造成麻烦。针对这两个问题，我们可以通过提高垃圾分拣的效率或者从源头来控制人们对于垃圾分类的错误率。

随着人工智能的不断发展，深度学习（Deep learning）在图像识别方面的技术已经较为成熟，例如在人脸识别、医学图像识别、遥感图像分类、ImageNet 分类识别、交通识别和字符识别等领域已取得了不错的成果[2]。深度学习可以对图像特征进行提取，通过大量训练来达到对图像进行分类和识别的功能[3]，因此，深度学习是图像处理的一个非常有效的技术。但深度学习在垃圾分类这一块的研究内容却不多，我们可以通过深度学习来改善人工分拣效率低、耗时耗力的问题，也可以通过深度学习来从源头上减少人们对垃圾识别的错误率，从而减少未处理垃圾的数量，让更多的资源可以被回收利用。针对当前城市垃圾分类现状，如果使用人工智能对垃圾进行识别和分类，不仅可以提高垃圾分拣厂的效率，还能够增加资源的回收利用率，改善垃圾对环境的污染程度，所以在环境保护和资源回收利用方面具有重大意义。

值得注意的是，垃圾有类间相似、类内差异大的特点。如何针对这些特点设计出适合垃圾分类的神经网络模型并达到预期的效果是智能垃圾识别的一个挑战。

2 基于YOLOv4的算法优化

2.1 YOLOv4算法介绍

YOLOv4（You Only Look Once, Version 4）模型是Bochkovskiy于2020年在YOLOv3[4]的基础上提出的。与YOLOv3网络结构相比，有如下几个改进：

（1）在输入端采用Mosaic 数据增强方法，源自CutMix 增强[5]，通过随机缩放、裁剪和排列四张图片来增加目标数量并增强背景的复杂度，极大地丰富了检测数据集，增加了小目标的数量，提高了网络的鲁棒性。

（2）在主网部分，YOLOv4 基于YOLOv3 主网Darknet53 和CSPNet 构建CSPDarknet53 网络。CSPDarknet53 网络不仅增强了CNN 网络的学习能力，而且在保证准确率的同时，减少了计算量，降低了内存成本。

（3）为了更好地提取目标的融合特征，YOLOv4 网络在主网络和输出层之间插入了SPP（空间金字塔池化）模块和FPN+PAN。首先，通过使用SPP 模块，可以放大主干特征接受范围，并且可以显着分离最重要的上下文特征。然后使用FPN 结构从上到下传递强大的语义特征，最后使用PAN 模块构建的特征金字塔结构从下到上传递强大的定位特征。通过该方法，可以融合不同层之间的特征。FPN（区域候选网络）+PAN（路径聚合网络）结构聚合了不同主干层的不同检测层的参数，进一步提高了网络的特征提取能力。

（4）使用Mish 激活函数替换主干中的Leaky ReLU 进行CSPDarknet53 分类器训练，使用分类器预训练权重可以提高分类器和检测器的准确率和泛化能力。

（5）Head 部分继承了YOLOv3 的Head 结构，共有三个分支。预测信息包括对象坐标、类别和置信度。

（6）使用了使用完全交叉并集（CIOU）损失代替IOU_loss作为边界框区域损失，还使用DIOU_nms 代替NMS，使得预测框回归的速度和精度更高。

（7）在由一系列下采样和多个残差网络组成的DarketNet53 网络中使用resblock_body 模块。

YOLOv4模型分为三个部分：CSPDarketnet53、Neck 和Head。CSPDarketnet53 结构是模型的骨干网络。Neck 结构由SPP和PAN 组成，Head 是模型的预测部分。YOLOv4 神经网络的结构如图1所示。可以看出，与YOLOv3 一样，在YOLOv4 中，输入图像被处理为宽度和高度为416×416 的图像。对于不同大小的目标， YOLO 网络的最终在YOLO Head 之后的输出层有宽和高三种尺寸，即13×13、26×26 和52×52。注意在输出层中，例如13×13×(c+5)×3，c 表示要检测的种类数；5 反映了anchor box 中是否有目标以及anchor box 的4 个调整参数；3 表示对应于每个宽度和高度大小的三个不同的anchor box。YOLOv4 网络的基本卷积块，即CBL 和CBM，与YOLOv3 类似，主要区别在于YOLOv4在骨干网络中采用Mish 函数作为激活函数，而在YOLOv3 中使用Leaky ReLU 函数。

图1：YOLOv4 网络结构图

YOLOv4 优点：

（1）增强CNN 的学习能力，使得在轻量化的同时保持准确性，速度和精度并存；

（2）降低内存成本；

（3）去掉了计算力较高的计算瓶颈结构，降低计算；

（4）是一种高效而强大的目标检测模型，使用 1080Ti 或2080Ti 就能训练出超快、准确的目标检测器；

（5）更适合单GPU 训练。

YOLOv4 缺点：

（1）在精度和速度上还有可以提升的空间；

（2）因为没有进行区域采样，所以对全局信息有较好的表现，但是在小范围的信息上表现较差。

由文献[6]可知，YOLOv4 在速度和精度上已超过Faster R-CNN和SSD，并且将YOLOv3 的AP 和FPS 分别提高10%和12%，也比较符合垃圾检测的条件，因此，本文将使用YOLOv4 算法并根据垃圾的特性对其进行改进，来完成垃圾检测。

2.2 基于注意力机制的YOLOv4算法改进

2.2.1 Triplet Attention

近年来，注意力机制已被应用于计算机视觉的各种任务。Non-local[7]通过计算不同位置的特征之间的依赖关系来捕获远程特征。SENet 提出了通道注意，它明确地建模了通道之间的依赖关系。GCNet[8]优化了Non-local，结合了全局注意力和通道注意力。SKNet[9]使用通道注意力机制来实现场感知的动态选择。卷积块注意模块(CBAM)在序列模式下结合了通道和空间注意机制。

传统的通道注意力计算会选取高层特征图进行全局池化，以具有全局的感受野，但这种计算方式会丢失大量的空间信息，也无法体现通道维数和空间维数之间的相互依赖性，传统的空间注意力计算同理。基于Spatial 和Channel 的CBAM 虽然可以体现通道与空间之间的相互依赖性，但通道注意与空间注意的计算是分离的，且在混合维度进行添加会带来非常巨大的计算量的消耗。

Triplet attention[10]相对于其它注意力机制有以下优点：

（1）可以保留大量空间和通道信息，避免大量丢失；

（2）能够进行跨维度的交互，强调了多维交互而不降低维度的重要性，因此消除了通道和权重之间的间接对应；

（3）只需要极小的算力，几乎没有参数。

我们选择在主干网络的最后一层添加Triplet attention。融合空间和通道注意力的注意力模块通常分别对空间和通道依赖进行建模，这导致特征的不同维度之间缺乏语义交互。Triplet attention 提取了不同维度之间的语义依赖，消除了channels 和weights 之间的间接对应关系，在不涉及到维数降低的情况下以很小的计算开销达到提高准确率的效果，廉价有效。图2显示了Triplet attention 的基本结构。顾名思义，triplet attention 由三个并行分支组成，其中两个分支负责捕获通道维度和空间维度之间的跨维度交互。剩下的最后一个分支类似于CBAM，用于构建空间注意力。所有三个分支的输出都使用简单的平均法聚合来实现跨维度的交互。在前两个分支中，triplet attention 分别沿H 轴和W 轴逆时针旋转原始输入张量90°，并将张量的形状从C×H×W 变换为W×H×C 和H×C×W。在第三个分支中，张量以其原始形状C×H×W 输入。之后，第一维的张量通过Z-pool 层降到第二维，得到平均聚合特征和最大聚合特征相连。这里的Z-pool 层负责通过连接该维度上的平均池化和最大池化特征，将张量的第零维减少到两个维度。这使该层能够保留实际张量的丰富表示，同时缩小其深度，使进一步的计算变得轻量级。Z-pool 定义为：

图2：Triplet attention

其中0d 是发生最大和平均池化操作的第0 维。例如，形状为(C×H×W)的张量的Z-Pool 将得到形状为(2×H×W)的张量。

然后，将缩减后的张量通过核大小为K 的标准卷积层、批量归一化层，最后将sigmoid 函数生成的对应维度的注意力权重加入到旋转张量中。 在最终输出时，第一个分支的输出沿H 轴顺时针旋转90 度，第二个分支的输出沿W 轴顺时针旋转90 度，确保形状与输入相同。最后，将三个分支的输出平均聚合为输出。综上所述，对于输入张量x∈RC×H×W，从Triplet attention 获得精炼的注意力应用张量y 的过程可以用以下等式表示：

其中σ 表示sigmoid 激活函数；ψ1、ψ2和ψ3分别代表了Triplet attention 的三个分支中由核大小k 定义的标准二维卷积层。简化式（3），y 变为：

其中ω1、ω2和ω3是三重注意力中计算的三个跨维度注意力权重。等式中的y1和y2表示顺时针旋转90°，以保留(C×H×W)的原始输入形状。

2.2.2 改进后的CSPDarknet53

本文将Triplet attention 加入YOLOv4 主干网络CSPDarknet53中来改进主干网络，这样可以使网络在几乎不添加参数的情况下能够自动学习进行跨维度的交互，提升有效的特征通道权重，从而让网络重点关注重要的特征通道，加强特征提取能力。改进后的CSPDarknet53 主干网络如图3所示。我们在每一组CSPResblock后面加一个注意力机制，由于注意力机制对特征图的尺寸和通道数都不会发生变化，所以对下一模块的输入不会产生影响。

图3：改进后的CSPDarknet53 网络

3 实验及结果分析

3.1 数据集

目前关于垃圾分类公开数据集主要有两个，TrashNet 数据集和华为垃圾数据集。还有一部分自制数据集，如医疗废弃物、街道垃圾等。TrashNet 数据集从分类情况上来看，主要是对可回收垃圾进行分类，并且缺少了布料这一类别，又由于本文主要针对上海垃圾分类标准进行分类识别，所以主要选取华为垃圾数据集作为参考。但华为数据集主要针对图像识别进行分类，没有对图像进行标注，所以在网络上搜索到其它研究者标注好的在飞桨平台AI Studio 公开的数据集。除此之外，还通过网络下载和拍照的方式对数据集进行了补充。因此，本文的数据集主要由以下三种形式组合而成：

（1）飞桨平台AI Studio 公开的标注好的华为垃圾数据集，并对其进行了筛选和调整。

（2）网络上搜索并通过筛选的相关垃圾图片。

（3）对日常生活垃圾进行拍照得到的垃圾图片。

本实验中垃圾数据集一共包含14964 张垃圾图片，分为4 大类，44 小类。由于目标检测的标签一般为英文，每一小类后面的英文为各类别对应的英文名，格式为“大类代号+小类类别名称”，其中1 为其它垃圾，2 为厨余垃圾，3 为可回收垃圾，4 为有害垃圾。由于图片中可能不止包含一个目标，会有多个标注，所以标注数会大于收集到的图片数，最终得到24248 个标注。如表1所示，其中有5467 个其它垃圾，4234 个厨余垃圾，12874 个可回收垃圾和1673 个有害垃圾。图4为每一小类的标注数量，最多的为陶瓷器皿1263 个，最少为金属厨具285 个。

图4：垃圾数据集子集

表1：垃圾数据集

在垃圾数据集中，由于大类中还分为了各个小类，会导致同一个大类中小类和小类之间差距较小，例如：饮料瓶和酒瓶外观相似，需要从更细节的特征来区分。而由于每个小类中也有不同尺度、也可能由于外力对形状产生形变，例如：被挤压的易拉罐、揉成一团的衣服，各小类垃圾之间的差异较大。因此，类内差距大，类间差异小，导致出现大量的错误检测。为了解决这些困难的样本，我们引入了注意力机制，以优化网络的特征提取能力，有效提高检测精度。

3.2 实验环境

本实验的硬件环境：处理器为Inter(R) Core(TM) i7-4710HQ CPU @ 2.50GHz，内存为12GB，显卡为GeForce GTX 1080Ti，操作系统是64 位，磁盘大小1T。

软件环境：操作系统为Windows 10，深度学习框架为pyTorch 1.9.0，开发环境为Visual Studio Code，开发语言为Python，包管理器和环境管理器为Anaconda，GPU 为CUDA 11.1 和cuDNN 8.0.5。

3.3 超参数设置

本实验将采用YOLOv4 在COCO 数据集上训练出来的预训练模型进行迁移学习，将预训练权重用于垃圾检测实验中进行训练，输入图像尺寸为416×416。垃圾识别的使用迁移学习的流程图如图5所示。

图5：垃圾识别任务流程图

本实验的学习率设置采取了余弦退火衰减（Cosine Annealing）策略，余弦退火可以通过余弦函数来降低学习率。我们将初始的学习率设为0.001，最小学习率设为0.00001，总共训练100 个epoch。本实验采用Adam 对整个网络进行迭代优化，我们将动量设为0.9，衰减因子设为0.0002，mini-batch 设为8。

3.4 实验结果

本节实验在CSPDarknet53 中加入注意力机制Triplet Attention，在实验参数设置统一的情况下，将YOLOv4 网络和改进后的网络在上一章提到的垃圾数据集的训练集上进行训练和验证，并且在测试集上进行测试和评估。在实验过程中统计了mAP（平均精度的平均值）和FPS（速度）用来作为本实验主要的评价指标。原YOLOv4 的实验结果和加入注意力机制后的实验结果如表2所示。

表2：注意力机制实验对比结果

由测试结果可知，加入Triplet Attention 注意力机制后虽然速度从28 降到26，但可以将mAP 从56.66%提升到58.57%，提升了1.91%。因此，我们可以得出加入Triplet Attention 到CSPDarknet53能够在几乎不增加参数和牺牲少量速度的条件下显著提高模型的性能。

为了更直观地展示检测效果，YOLOv4 和改进后的检测结果如图6所示。左图为原YOLOv4，右图为改进后的算法，可以看出，改进后的YOLOv4 对上述测试图像有更好的检测效果。

图6：检测效果对比图

4 结论

本文主要对YOLOv4 在垃圾分类上的应用展开了研究，在研究的基础上，针对垃圾数据集的特点，对YOLOv4 网络模型做出改进。由于垃圾分类数据集同一个大类中小类和小类之间差距较小，而由于每个小类中也有不同尺度、也可能由于外力对形状产生形变，各小类垃圾之间的差异较大，需要从更细节的特征来区分。我们将注意力机制Triplet attention 加入YOLOv4 主干网络CSPDarknet53中来改进主干网络，这样可以使网络在几乎不添加参数的情况下能够自动学习进行跨维度的交互，提升有效的特征通道权重，从而让网络重点关注重要的特征通道，提高网络检测精度。最后，基于以上改进，本文在自己制作的垃圾数据集上将原YOLOv4 算法与改进后的算法进行对比实验，实验结果证明本文的改进对垃圾检测精度方面有所提升，说明此改进对垃圾分类研究是有效的，对智能垃圾分类技术提供了参考。