基于遷移學(xué)習(xí)的垃圾圖像分類模型的比較研究

牛鐿潼，孟德堯，吳榮海，任國鵬，楊鄧奇*

（1.大理大學(xué)數(shù)學(xué)與計(jì)算機(jī)學(xué)院，云南大理 671003；2.大理大學(xué)工程實(shí)訓(xùn)中心，云南大理 671003；3.大理大學(xué)農(nóng)學(xué)與生物科學(xué)學(xué)院，云南大理 671003）

隨著經(jīng)濟(jì)的不斷發(fā)展，居民生活水平不斷提高，生活垃圾的產(chǎn)生量也越來越多，如果垃圾沒有得到妥善處理，將會使環(huán)境問題變得更加嚴(yán)峻〔1〕。垃圾的分類處理、回收利用是垃圾處理的必然選擇。當(dāng)前的垃圾分類主要靠人工完成，存在效率低、成本高等問題。利用計(jì)算機(jī)視覺技術(shù)實(shí)現(xiàn)垃圾自動分類，可以為下一步機(jī)器人自動分揀垃圾提供技術(shù)支持。機(jī)器人通過自動識別技術(shù)在垃圾傳送帶上完成垃圾自動分揀。利用計(jì)算機(jī)技術(shù)實(shí)現(xiàn)垃圾自動分揀節(jié)約人力資源的同時也能提高垃圾處理效率。深度學(xué)習(xí)技術(shù)在圖像識別領(lǐng)域應(yīng)用廣泛，已有學(xué)者將深度學(xué)習(xí)技術(shù)應(yīng)用于垃圾分類，取得了突出的成績，例如Rad等〔2〕在車輛上安裝了圖像采集系統(tǒng)，采用基于深度學(xué)習(xí)的框架對不同類型的垃圾進(jìn)行定位和分類；Bobulski等〔3〕提出了一個使用卷積神經(jīng)網(wǎng)絡(luò)對塑料垃圾進(jìn)行分類的系統(tǒng)，將垃圾分為4個類別，實(shí)驗(yàn)證明了分類系統(tǒng)的有效性。

現(xiàn)有關(guān)于垃圾自動識別的研究主要以包括紙張、玻璃、塑料、金屬、紙板和其他垃圾6個類別的TrashNet〔4〕標(biāo)桿數(shù)據(jù)集為基礎(chǔ)，利用不同深度學(xué)習(xí)模型實(shí)現(xiàn)垃圾分類。例如，Japkowicz等〔5〕微調(diào)了一種輕量級的深層神經(jīng)網(wǎng)絡(luò)MobileNet，在TrashNet數(shù)據(jù)集上獲得87.2%的準(zhǔn)確率；Aral等〔6〕又微調(diào)了DenseNet、Inception-v4、Xception以及MobileNet網(wǎng)絡(luò)，在TrashNet數(shù)據(jù)集上進(jìn)行訓(xùn)練及測試，得出使用微調(diào)后的DenseNet網(wǎng)絡(luò)時效果最好，達(dá)到95%的準(zhǔn)確率。TrashNet數(shù)據(jù)集具有背景簡單、特征突出等特點(diǎn)，因此現(xiàn)有基于TrashNet數(shù)據(jù)集的方法往往都取得較高的準(zhǔn)確率。

當(dāng)前的深度學(xué)習(xí)模型包括傳統(tǒng)的深度卷積神經(jīng)網(wǎng)絡(luò)和基于注意力機(jī)制的深度神經(jīng)網(wǎng)絡(luò)?，F(xiàn)有的研究幾乎都使用傳統(tǒng)的深度卷積神經(jīng)網(wǎng)絡(luò)，使用基于注意力機(jī)制的深度學(xué)習(xí)模型進(jìn)行垃圾分類尚無可查的文獻(xiàn)。哪種類型的深度神經(jīng)網(wǎng)絡(luò)在小規(guī)模垃圾分類數(shù)據(jù)集上表現(xiàn)更好是個值得探討的問題。本文選取3個具有代表性的深度卷積神經(jīng)網(wǎng)絡(luò)模型（VGG16、ResNet50和ResNeXt50）和2個注意力機(jī)制深度神經(jīng)網(wǎng)絡(luò)模型（ViT和ViT-Hybrid）進(jìn)行一系列實(shí)驗(yàn)，系統(tǒng)地比較這2種類型的神經(jīng)網(wǎng)絡(luò)在小規(guī)模垃圾分類數(shù)據(jù)集上的性能，為其他學(xué)者在垃圾分類模型選擇和性能評估方面提供參考。

1 研究現(xiàn)狀

在我國大力提倡垃圾分類回收、保護(hù)環(huán)境的大背景下，許多學(xué)者也在這一領(lǐng)域開展了研究。彭韻等〔7〕研究了我國垃圾分類的發(fā)展歷程，分析了垃圾分類過程中存在的問題和障礙，并為垃圾分類的管理提出了對策和建議；Han等〔8〕研究了影響發(fā)展中國家農(nóng)村地區(qū)生活垃圾產(chǎn)生的因素，包括社會因素、經(jīng)濟(jì)因素和自然因素；Duan等〔9〕分析了我國有代表性的城市生活垃圾的產(chǎn)生情況及其影響因素，并利用神經(jīng)網(wǎng)絡(luò)預(yù)測了城市垃圾的產(chǎn)生情況。在現(xiàn)實(shí)生活中大多數(shù)人對垃圾分類內(nèi)容的認(rèn)知都停留在基礎(chǔ)階段，對分類標(biāo)準(zhǔn)和如何操作都不是很清楚，而且垃圾的種類也比較多，想完全正確分類也是有一定難度的，這也表明，如果想更好地實(shí)施垃圾分類，僅靠民眾的認(rèn)知是遠(yuǎn)遠(yuǎn)不夠的，需要相關(guān)的計(jì)算機(jī)技術(shù)來幫助人們確定垃圾的類別。

目前，垃圾的分類識別在技術(shù)方面主要是以基于卷積的深度學(xué)習(xí)模型為主，基于注意力機(jī)制的神經(jīng)網(wǎng)絡(luò)模型在垃圾分類方面的應(yīng)用尚無可查的文獻(xiàn)。利用帶標(biāo)簽的訓(xùn)練樣本建立深度神經(jīng)網(wǎng)絡(luò)模型，再用建好的模型完成垃圾自動識別與分類。Zhang等〔10〕提出了一種基于遷移學(xué)習(xí)的Dense-Net169垃圾圖像分類模型，并自制了一個垃圾圖像數(shù)據(jù)集，實(shí)驗(yàn)結(jié)果表明，DenseNet169模型經(jīng)過遷移學(xué)習(xí)后的分類準(zhǔn)確率在82%以上，優(yōu)于其他圖像分類算法。Sai Susanth等〔11〕將AlexNet、VGG16、ResNet50、DenseNet169在數(shù)據(jù)集TrashNet上進(jìn)行驗(yàn)證，結(jié)果表明，DenseNet169性能更好，檢測準(zhǔn)確率達(dá)到94.9%。Mao等〔12〕利用TrashNet數(shù)據(jù)集上優(yōu)化的卷積神經(jīng)網(wǎng)絡(luò)對垃圾進(jìn)行分類，利用遺傳算法優(yōu)化DenseNet121的全連接層，優(yōu)化后的模型在TrashNet數(shù)據(jù)集上達(dá)到了99.6%的分類準(zhǔn)確率。Song等〔13〕提出基于ResNet和Inception-v4的垃圾自動分類DSCR網(wǎng)絡(luò)，準(zhǔn)確率達(dá)到94.38%。Wang等〔14〕提出的自動垃圾檢測系統(tǒng)使用ResNet網(wǎng)絡(luò)作為卷積層，輸入數(shù)據(jù)集圖像后通過ResNet的共享卷積層生成特征圖，該方法很好地克服了垃圾檢測時產(chǎn)生的區(qū)域誤檢問題。Gu等〔15〕基于ResNet50設(shè)計(jì)了一種輕量級垃圾分類模型，利用深度卷積和分組卷積減少計(jì)算量和參數(shù)，相對于原始模型，該模型分類準(zhǔn)確率更高，占用內(nèi)存更少，能夠解決垃圾分類在嵌入式設(shè)備上的應(yīng)用問題。Vo等〔16〕修改了原始ResNeXt模型，在全局平均池化層之后添加了2個全連接層以提高網(wǎng)絡(luò)的分類性能，能夠快速實(shí)現(xiàn)分類模型的穩(wěn)定和泛化，相比于ResNet101，該方法對金屬、紙張和塑料的分類性能表現(xiàn)較差，如何利用這些樣本的獨(dú)有特征提高分類準(zhǔn)確率值得深入探索。Yang等〔17〕基于ResNeXt101設(shè)計(jì)了一個新的增量學(xué)習(xí)框架，以解決垃圾分類面臨的缺乏足夠數(shù)據(jù)、高成本的類別增量和噪聲標(biāo)簽等問題，使用增量學(xué)習(xí)方法使模型不斷從新樣本中學(xué)習(xí)和更新，在華為云垃圾數(shù)據(jù)集上以可接受的推理速度達(dá)到了96.96%。

2 材料和方法

2.1 數(shù)據(jù)集

2.1.1 數(shù)據(jù)集描述 本文使用了公開的TrashNet和自制的NewTrash 2個垃圾分類數(shù)據(jù)集。TrashNet是垃圾分類標(biāo)桿數(shù)據(jù)集，目前，基于垃圾圖像識別分類的研究多以TrashNet數(shù)據(jù)集為主，它是由Thung和Yang于2016年創(chuàng)建的數(shù)據(jù)集，由2 527張圖像組成，包含6類垃圾的RGB圖像，分別是紙張（594張）、玻璃（501張）、塑料（482張）、金屬（410張）、紙板（403張）和其他垃圾（137張）。TrashNet數(shù)據(jù)集制作是通過將物體放置在白色海報(bào)板上并使用陽光或室內(nèi)照明來捕獲圖像，這樣的垃圾圖像背景是單一的，前景目標(biāo)特征明顯。模型訓(xùn)練和識別往往可以獲得很高的準(zhǔn)確率。

自制的NewTrash垃圾分類數(shù)據(jù)集共3 281張垃圾圖像，包括7個類別，分別是紙（522張）、玻璃（501張）、塑料（551張）、金屬（551張）、藥物（448張）、軟膏（387張）和電池（321張）。相比TrashNet數(shù)據(jù)集，NewTrash數(shù)據(jù)集中圖像背景相對復(fù)雜，并將普通的紙張和硬紙板合并為一類，同時增加了藥物、電池和軟膏。NewTrash也是一個小數(shù)據(jù)集，但是背景比TrashNet復(fù)雜化了且類別和數(shù)量比其豐富。見圖1。

圖1 NewTrash數(shù)據(jù)集圖像示例

2.1.2 數(shù)據(jù)集劃分 本文基于這2個垃圾數(shù)據(jù)集進(jìn)行模型分類準(zhǔn)確率和泛化能力的實(shí)驗(yàn)，2個數(shù)據(jù)集的劃分都采用8∶1∶1的比例，80%作為訓(xùn)練集，10%作為驗(yàn)證集，10%作為測試集，數(shù)據(jù)集中各類別詳細(xì)劃分見表1～2。

表1 TrashNet數(shù)據(jù)集劃分

表2 NewTrash數(shù)據(jù)集劃分

2.2 模型性能指標(biāo) 通過模型分類性能指標(biāo)來評價垃圾分類模型的性能，本文采用準(zhǔn)確率對垃圾分類模型性能進(jìn)行評價。準(zhǔn)確率表示垃圾分類模型分類正確的比例，定義如公式（1）所示：

其中N為樣本總數(shù)，n為樣本總類別，Xii為第i類樣本分類正確的樣本數(shù)，一個模型中所有樣本分類正確的總和與樣本總數(shù)的比值即為模型的分類準(zhǔn)確率。

2.3 模型性能比較的方案設(shè)計(jì) 在深度神經(jīng)網(wǎng)絡(luò)中有諸多較好的分類網(wǎng)絡(luò)，本研究選用比較有代表性的5個深度神經(jīng)網(wǎng)絡(luò)作為垃圾分類模型，分別是VGG16、ResNet50、ResNeXt50、ViT和ViT-Hybrid 5個模型。見表3。

表3 所用網(wǎng)絡(luò)簡介

本文訓(xùn)練模型時，使用了遷移學(xué)習(xí)，將在ImageNet上訓(xùn)練好的網(wǎng)絡(luò)權(quán)重參數(shù)遷移到2個垃圾數(shù)據(jù)集上，并對模型的全連接層進(jìn)行微調(diào)。為了比較5個模型在單一數(shù)據(jù)集上的性能和模型的泛化性能，本文設(shè)計(jì)了2個實(shí)驗(yàn)方案。

2.3.1 單一數(shù)據(jù)集上模型性能比較實(shí)驗(yàn)設(shè)計(jì) 為了測試模型的分類性能，本文設(shè)計(jì)了在單一數(shù)據(jù)集上測試的方案，即分別在TrashNet和NewTrash數(shù)據(jù)集上訓(xùn)練和測試各個模型的分類準(zhǔn)確率。如圖2所示，對于TrashNet數(shù)據(jù)集，將ImageNet上預(yù)訓(xùn)練好的模型分別遷移到TrashNet的訓(xùn)練集上進(jìn)行訓(xùn)練，用TrashNet上的驗(yàn)證集監(jiān)測模型的訓(xùn)練過程，訓(xùn)練完成后的5個模型分別在TrashNet上的測試集測試。對于NewTrash數(shù)據(jù)集，將ImageNet上預(yù)訓(xùn)練好的5個模型分別遷移到NewTrash的訓(xùn)練集上進(jìn)行訓(xùn)練，用NewTrash上的驗(yàn)證集監(jiān)測模型訓(xùn)練過程，訓(xùn)練完成后的5個模型分別在NewTrash上的測試集測試。

圖2 單一數(shù)據(jù)集上模型性能比較實(shí)驗(yàn)設(shè)計(jì)

2.3.2 模型泛化性能比較實(shí)驗(yàn)設(shè)計(jì) 在分類模型中，模型分類準(zhǔn)確率的高低不代表模型的整體性能好壞，模型的性能還可以從模型泛化能力的高低來判斷，泛化能力是指機(jī)器學(xué)習(xí)算法對新鮮樣本的適應(yīng)能力，也就是用訓(xùn)練好的模型去識別樣本外數(shù)據(jù)（不在訓(xùn)練集中的圖像），如果能在2個完全獨(dú)立的數(shù)據(jù)集上實(shí)驗(yàn)，能更真實(shí)地反映模型的泛化性。為了驗(yàn)證實(shí)驗(yàn)所選模型的泛化能力，本文基于TrashNet和NewTrash 2個數(shù)據(jù)集，設(shè)計(jì)了模型泛化性能比較實(shí)驗(yàn)。5個ImageNet上預(yù)訓(xùn)練的模型分別遷移到TrashNet的訓(xùn)練集上進(jìn)行訓(xùn)練，用TrashNet的驗(yàn)證集監(jiān)測模型的訓(xùn)練過程，最后用完整的NewTrash數(shù)據(jù)集進(jìn)行測試。同樣地，5個ImageNet預(yù)訓(xùn)練的模型分別遷移到在NewTrash的訓(xùn)練集上，用NewTrash的驗(yàn)證集監(jiān)測模型的訓(xùn)練過程，最后用完整的TrashNet數(shù)據(jù)集進(jìn)行測試。見圖3。

圖3 模型泛化性能比較實(shí)驗(yàn)設(shè)計(jì)

由于TrashNet和NewTrash 2個數(shù)據(jù)集包含的類別不完全一致，僅有玻璃、金屬和塑料3個類別是2個數(shù)據(jù)集共有的。所以，這2個數(shù)據(jù)集用于模型泛化能力測試時，被劃分為4個類別，即共有的3類和其他1個類別。在模型測試過程中分類正確的情況是模型將玻璃、金屬和塑料分別正確地識別為玻璃、金屬和塑料，將其他類別識別為其他，即真實(shí)標(biāo)簽與預(yù)測標(biāo)簽一致時則為分類正確。見圖4。

圖4 模型測試時分類識別正確的情況

3 實(shí)驗(yàn)結(jié)果與分析

VGG16、ResNet50、ResNeXt50、ViT、ViT-Hybrid的訓(xùn)練和測試在PyTorch深度學(xué)習(xí)框架下運(yùn)行。訓(xùn)練模型的服務(wù)器配置為Intel（R）Xeon（R）CPU E5-2650 v3@2.30 GHz，NVIDIA Quadro M5000，Ubuntu18作為操作系統(tǒng)。訓(xùn)練模型時都使用SGD優(yōu)化器，BatchSize設(shè)為20，Epoch設(shè)為40。

單一數(shù)據(jù)集上模型性能測試的實(shí)驗(yàn)結(jié)果見圖5。在TrashNet數(shù)據(jù)集上使用的5個模型中，由分類準(zhǔn)確率來看基于注意力機(jī)制的深度神經(jīng)網(wǎng)絡(luò)模型ViT和ViT-Hybrid比基于卷積的深度神經(jīng)網(wǎng)絡(luò)模型VGG16、ResNet50和ResNeXt50的分類效果更好，且純注意力機(jī)制的深度神經(jīng)網(wǎng)絡(luò)模型ViT比卷積和注意力機(jī)制結(jié)合的神經(jīng)網(wǎng)絡(luò)模型ViT-Hybrid的分類效果更好。在NewTrash數(shù)據(jù)集上的分類效果也是基于注意力機(jī)制的模型比基于卷積的模型分類效果好，而且純注意力機(jī)制模型ViT比卷積和注意力機(jī)制結(jié)合的神經(jīng)網(wǎng)絡(luò)ViT-Hybrid的分類效果好。2個數(shù)據(jù)集在5個模型上的分類效果整體來說都不錯，而且準(zhǔn)確率值相差較小，但可以明顯看出在NewTrash數(shù)據(jù)集上得到的分類準(zhǔn)確率沒有在TrashNet數(shù)據(jù)集上得到的分類準(zhǔn)確率高，出現(xiàn)這種情況的原因是自制的NewTrash數(shù)據(jù)集相對于TrashNet數(shù)據(jù)集的垃圾類別和數(shù)量多，且圖像背景更復(fù)雜。從單一數(shù)據(jù)集上模型性能測試的實(shí)驗(yàn)數(shù)據(jù)結(jié)果來看，純注意力機(jī)制深度神經(jīng)網(wǎng)絡(luò)模型ViT在小規(guī)模垃圾分類數(shù)據(jù)集上的分類性能最好。

圖5 模型分類準(zhǔn)確率

為了更好地驗(yàn)證ViT模型的性能，本實(shí)驗(yàn)在跨數(shù)據(jù)集上進(jìn)行泛化能力性能的測試，再次證明了ViT模型的性能優(yōu)勢。模型泛化性能比較的實(shí)驗(yàn)結(jié)果見圖6。在NewTrash訓(xùn)練集上訓(xùn)練得到的模型的泛化能力（NewTrash測TrashNet）都比在TrashNet訓(xùn)練集上訓(xùn)練得到的模型泛化能力（TrashNet測NewTrash）好，而且基于注意力機(jī)制的2個深度神經(jīng)網(wǎng)絡(luò)模型的泛化準(zhǔn)確率都比基于卷積的3個深度神經(jīng)網(wǎng)絡(luò)模型的泛化準(zhǔn)確率高。在NewTrash訓(xùn)練集上訓(xùn)練得到的純注意力機(jī)制網(wǎng)絡(luò)模型ViT的泛化準(zhǔn)確率最大。泛化實(shí)驗(yàn)結(jié)果可以表明純注意力機(jī)制深度神經(jīng)網(wǎng)絡(luò)模型ViT在小規(guī)模垃圾分類數(shù)據(jù)集上的泛化性能最好。

圖6 模型泛化準(zhǔn)確率

模型的分類性能和泛化性能的實(shí)驗(yàn)結(jié)果綜合表明，盡管ViT-Hybrid模型結(jié)合了注意力機(jī)制和卷積2種主流的深度學(xué)習(xí)思路的優(yōu)點(diǎn)，但在訓(xùn)練樣本有限的情況下，其總體性能不如純注意力機(jī)制的ViT模型。

4 總結(jié)

為了驗(yàn)證基于卷積和基于注意力機(jī)制的2類深度神經(jīng)網(wǎng)絡(luò)在小規(guī)模垃圾數(shù)據(jù)集上的性能，本文選取了幾種有代表性的模型開展了一系列實(shí)驗(yàn)，實(shí)驗(yàn)結(jié)果表明，對現(xiàn)有的小規(guī)模垃圾數(shù)據(jù)集來說，純注意力機(jī)制的模型更具有優(yōu)勢。

本文研究內(nèi)容有限，只是在幾個具有代表性的網(wǎng)絡(luò)模型上進(jìn)行實(shí)驗(yàn)，將深度學(xué)習(xí)技術(shù)應(yīng)用于垃圾分類領(lǐng)域還是具有很大的探索價值，今后將繼續(xù)豐富該方向的實(shí)驗(yàn)內(nèi)容。深度學(xué)習(xí)如今發(fā)展迅猛，相信將來在垃圾分類領(lǐng)域會出現(xiàn)更好的可供選擇的分類模型。