基于迁移学习的垃圾图像分类模型的比较研究

牛镱潼，孟德尧，吴荣海，任国鹏，杨邓奇*

（1.大理大学数学与计算机学院，云南大理 671003；2.大理大学工程实训中心，云南大理 671003；3.大理大学农学与生物科学学院，云南大理 671003）

随着经济的不断发展，居民生活水平不断提高，生活垃圾的产生量也越来越多，如果垃圾没有得到妥善处理，将会使环境问题变得更加严峻〔1〕。垃圾的分类处理、回收利用是垃圾处理的必然选择。当前的垃圾分类主要靠人工完成，存在效率低、成本高等问题。利用计算机视觉技术实现垃圾自动分类，可以为下一步机器人自动分拣垃圾提供技术支持。机器人通过自动识别技术在垃圾传送带上完成垃圾自动分拣。利用计算机技术实现垃圾自动分拣节约人力资源的同时也能提高垃圾处理效率。深度学习技术在图像识别领域应用广泛，已有学者将深度学习技术应用于垃圾分类，取得了突出的成绩，例如Rad等〔2〕在车辆上安装了图像采集系统，采用基于深度学习的框架对不同类型的垃圾进行定位和分类；Bobulski等〔3〕提出了一个使用卷积神经网络对塑料垃圾进行分类的系统，将垃圾分为4个类别，实验证明了分类系统的有效性。

现有关于垃圾自动识别的研究主要以包括纸张、玻璃、塑料、金属、纸板和其他垃圾6个类别的TrashNet〔4〕标杆数据集为基础，利用不同深度学习模型实现垃圾分类。例如，Japkowicz等〔5〕微调了一种轻量级的深层神经网络MobileNet，在TrashNet数据集上获得87.2%的准确率；Aral等〔6〕又微调了DenseNet、Inception-v4、Xception以及MobileNet网络，在TrashNet数据集上进行训练及测试，得出使用微调后的DenseNet网络时效果最好，达到95%的准确率。TrashNet数据集具有背景简单、特征突出等特点，因此现有基于TrashNet数据集的方法往往都取得较高的准确率。

当前的深度学习模型包括传统的深度卷积神经网络和基于注意力机制的深度神经网络。现有的研究几乎都使用传统的深度卷积神经网络，使用基于注意力机制的深度学习模型进行垃圾分类尚无可查的文献。哪种类型的深度神经网络在小规模垃圾分类数据集上表现更好是个值得探讨的问题。本文选取3个具有代表性的深度卷积神经网络模型（VGG16、ResNet50和ResNeXt50）和2个注意力机制深度神经网络模型（ViT和ViT-Hybrid）进行一系列实验，系统地比较这2种类型的神经网络在小规模垃圾分类数据集上的性能，为其他学者在垃圾分类模型选择和性能评估方面提供参考。

1 研究现状

在我国大力提倡垃圾分类回收、保护环境的大背景下，许多学者也在这一领域开展了研究。彭韵等〔7〕研究了我国垃圾分类的发展历程，分析了垃圾分类过程中存在的问题和障碍，并为垃圾分类的管理提出了对策和建议；Han等〔8〕研究了影响发展中国家农村地区生活垃圾产生的因素，包括社会因素、经济因素和自然因素；Duan等〔9〕分析了我国有代表性的城市生活垃圾的产生情况及其影响因素，并利用神经网络预测了城市垃圾的产生情况。在现实生活中大多数人对垃圾分类内容的认知都停留在基础阶段，对分类标准和如何操作都不是很清楚，而且垃圾的种类也比较多，想完全正确分类也是有一定难度的，这也表明，如果想更好地实施垃圾分类，仅靠民众的认知是远远不够的，需要相关的计算机技术来帮助人们确定垃圾的类别。

目前，垃圾的分类识别在技术方面主要是以基于卷积的深度学习模型为主，基于注意力机制的神经网络模型在垃圾分类方面的应用尚无可查的文献。利用带标签的训练样本建立深度神经网络模型，再用建好的模型完成垃圾自动识别与分类。Zhang等〔10〕提出了一种基于迁移学习的Dense-Net169垃圾图像分类模型，并自制了一个垃圾图像数据集，实验结果表明，DenseNet169模型经过迁移学习后的分类准确率在82%以上，优于其他图像分类算法。Sai Susanth等〔11〕将AlexNet、VGG16、ResNet50、DenseNet169在数据集TrashNet上进行验证，结果表明，DenseNet169性能更好，检测准确率达到94.9%。Mao等〔12〕利用TrashNet数据集上优化的卷积神经网络对垃圾进行分类，利用遗传算法优化DenseNet121的全连接层，优化后的模型在TrashNet数据集上达到了99.6%的分类准确率。Song等〔13〕提出基于ResNet和Inception-v4的垃圾自动分类DSCR网络，准确率达到94.38%。Wang等〔14〕提出的自动垃圾检测系统使用ResNet网络作为卷积层，输入数据集图像后通过ResNet的共享卷积层生成特征图，该方法很好地克服了垃圾检测时产生的区域误检问题。Gu等〔15〕基于ResNet50设计了一种轻量级垃圾分类模型，利用深度卷积和分组卷积减少计算量和参数，相对于原始模型，该模型分类准确率更高，占用内存更少，能够解决垃圾分类在嵌入式设备上的应用问题。Vo等〔16〕修改了原始ResNeXt模型，在全局平均池化层之后添加了2个全连接层以提高网络的分类性能，能够快速实现分类模型的稳定和泛化，相比于ResNet101，该方法对金属、纸张和塑料的分类性能表现较差，如何利用这些样本的独有特征提高分类准确率值得深入探索。Yang等〔17〕基于ResNeXt101设计了一个新的增量学习框架，以解决垃圾分类面临的缺乏足够数据、高成本的类别增量和噪声标签等问题，使用增量学习方法使模型不断从新样本中学习和更新，在华为云垃圾数据集上以可接受的推理速度达到了96.96%。

2 材料和方法

2.1 数据集

2.1.1 数据集描述 本文使用了公开的TrashNet和自制的NewTrash 2个垃圾分类数据集。TrashNet是垃圾分类标杆数据集，目前，基于垃圾图像识别分类的研究多以TrashNet数据集为主，它是由Thung和Yang于2016年创建的数据集，由2 527张图像组成，包含6类垃圾的RGB图像，分别是纸张（594张）、玻璃（501张）、塑料（482张）、金属（410张）、纸板（403张）和其他垃圾（137张）。TrashNet数据集制作是通过将物体放置在白色海报板上并使用阳光或室内照明来捕获图像，这样的垃圾图像背景是单一的，前景目标特征明显。模型训练和识别往往可以获得很高的准确率。

自制的NewTrash垃圾分类数据集共3 281张垃圾图像，包括7个类别，分别是纸（522张）、玻璃（501张）、塑料（551张）、金属（551张）、药物（448张）、软膏（387张）和电池（321张）。相比TrashNet数据集，NewTrash数据集中图像背景相对复杂，并将普通的纸张和硬纸板合并为一类，同时增加了药物、电池和软膏。NewTrash也是一个小数据集，但是背景比TrashNet复杂化了且类别和数量比其丰富。见图1。

图1 NewTrash数据集图像示例

2.1.2 数据集划分 本文基于这2个垃圾数据集进行模型分类准确率和泛化能力的实验，2个数据集的划分都采用8∶1∶1的比例，80%作为训练集，10%作为验证集，10%作为测试集，数据集中各类别详细划分见表1～2。

表1 TrashNet数据集划分

表2 NewTrash数据集划分

2.2 模型性能指标 通过模型分类性能指标来评价垃圾分类模型的性能，本文采用准确率对垃圾分类模型性能进行评价。准确率表示垃圾分类模型分类正确的比例，定义如公式（1）所示：

其中N为样本总数，n为样本总类别，Xii为第i类样本分类正确的样本数，一个模型中所有样本分类正确的总和与样本总数的比值即为模型的分类准确率。

2.3 模型性能比较的方案设计 在深度神经网络中有诸多较好的分类网络，本研究选用比较有代表性的5个深度神经网络作为垃圾分类模型，分别是VGG16、ResNet50、ResNeXt50、ViT和ViT-Hybrid 5个模型。见表3。

表3 所用网络简介

本文训练模型时，使用了迁移学习，将在ImageNet上训练好的网络权重参数迁移到2个垃圾数据集上，并对模型的全连接层进行微调。为了比较5个模型在单一数据集上的性能和模型的泛化性能，本文设计了2个实验方案。

2.3.1 单一数据集上模型性能比较实验设计 为了测试模型的分类性能，本文设计了在单一数据集上测试的方案，即分别在TrashNet和NewTrash数据集上训练和测试各个模型的分类准确率。如图2所示，对于TrashNet数据集，将ImageNet上预训练好的模型分别迁移到TrashNet的训练集上进行训练，用TrashNet上的验证集监测模型的训练过程，训练完成后的5个模型分别在TrashNet上的测试集测试。对于NewTrash数据集，将ImageNet上预训练好的5个模型分别迁移到NewTrash的训练集上进行训练，用NewTrash上的验证集监测模型训练过程，训练完成后的5个模型分别在NewTrash上的测试集测试。

图2 单一数据集上模型性能比较实验设计

2.3.2 模型泛化性能比较实验设计 在分类模型中，模型分类准确率的高低不代表模型的整体性能好坏，模型的性能还可以从模型泛化能力的高低来判断，泛化能力是指机器学习算法对新鲜样本的适应能力，也就是用训练好的模型去识别样本外数据（不在训练集中的图像），如果能在2个完全独立的数据集上实验，能更真实地反映模型的泛化性。为了验证实验所选模型的泛化能力，本文基于TrashNet和NewTrash 2个数据集，设计了模型泛化性能比较实验。5个ImageNet上预训练的模型分别迁移到TrashNet的训练集上进行训练，用TrashNet的验证集监测模型的训练过程，最后用完整的NewTrash数据集进行测试。同样地，5个ImageNet预训练的模型分别迁移到在NewTrash的训练集上，用NewTrash的验证集监测模型的训练过程，最后用完整的TrashNet数据集进行测试。见图3。

图3 模型泛化性能比较实验设计

由于TrashNet和NewTrash 2个数据集包含的类别不完全一致，仅有玻璃、金属和塑料3个类别是2个数据集共有的。所以，这2个数据集用于模型泛化能力测试时，被划分为4个类别，即共有的3类和其他1个类别。在模型测试过程中分类正确的情况是模型将玻璃、金属和塑料分别正确地识别为玻璃、金属和塑料，将其他类别识别为其他，即真实标签与预测标签一致时则为分类正确。见图4。

图4 模型测试时分类识别正确的情况

3 实验结果与分析

VGG16、ResNet50、ResNeXt50、ViT、ViT-Hybrid的训练和测试在PyTorch深度学习框架下运行。训练模型的服务器配置为Intel（R）Xeon（R）CPU E5-2650 v3@2.30 GHz，NVIDIA Quadro M5000，Ubuntu18作为操作系统。训练模型时都使用SGD优化器，BatchSize设为20，Epoch设为40。

单一数据集上模型性能测试的实验结果见图5。在TrashNet数据集上使用的5个模型中，由分类准确率来看基于注意力机制的深度神经网络模型ViT和ViT-Hybrid比基于卷积的深度神经网络模型VGG16、ResNet50和ResNeXt50的分类效果更好，且纯注意力机制的深度神经网络模型ViT比卷积和注意力机制结合的神经网络模型ViT-Hybrid的分类效果更好。在NewTrash数据集上的分类效果也是基于注意力机制的模型比基于卷积的模型分类效果好，而且纯注意力机制模型ViT比卷积和注意力机制结合的神经网络ViT-Hybrid的分类效果好。2个数据集在5个模型上的分类效果整体来说都不错，而且准确率值相差较小，但可以明显看出在NewTrash数据集上得到的分类准确率没有在TrashNet数据集上得到的分类准确率高，出现这种情况的原因是自制的NewTrash数据集相对于TrashNet数据集的垃圾类别和数量多，且图像背景更复杂。从单一数据集上模型性能测试的实验数据结果来看，纯注意力机制深度神经网络模型ViT在小规模垃圾分类数据集上的分类性能最好。

图5 模型分类准确率

为了更好地验证ViT模型的性能，本实验在跨数据集上进行泛化能力性能的测试，再次证明了ViT模型的性能优势。模型泛化性能比较的实验结果见图6。在NewTrash训练集上训练得到的模型的泛化能力（NewTrash测TrashNet）都比在TrashNet训练集上训练得到的模型泛化能力（TrashNet测NewTrash）好，而且基于注意力机制的2个深度神经网络模型的泛化准确率都比基于卷积的3个深度神经网络模型的泛化准确率高。在NewTrash训练集上训练得到的纯注意力机制网络模型ViT的泛化准确率最大。泛化实验结果可以表明纯注意力机制深度神经网络模型ViT在小规模垃圾分类数据集上的泛化性能最好。

图6 模型泛化准确率

模型的分类性能和泛化性能的实验结果综合表明，尽管ViT-Hybrid模型结合了注意力机制和卷积2种主流的深度学习思路的优点，但在训练样本有限的情况下，其总体性能不如纯注意力机制的ViT模型。

4 总结

为了验证基于卷积和基于注意力机制的2类深度神经网络在小规模垃圾数据集上的性能，本文选取了几种有代表性的模型开展了一系列实验，实验结果表明，对现有的小规模垃圾数据集来说，纯注意力机制的模型更具有优势。

本文研究内容有限，只是在几个具有代表性的网络模型上进行实验，将深度学习技术应用于垃圾分类领域还是具有很大的探索价值，今后将继续丰富该方向的实验内容。深度学习如今发展迅猛，相信将来在垃圾分类领域会出现更好的可供选择的分类模型。