基于遷移學習的垃圾圖像分類模型的比較研究

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

（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類深度神經網絡在小規模垃圾數據集上的性能，本文選取了幾種有代表性的模型開展了一系列實驗，實驗結果表明，對現有的小規模垃圾數據集來說，純注意力機制的模型更具有優勢。

本文研究內容有限，只是在幾個具有代表性的網絡模型上進行實驗，將深度學習技術應用于垃圾分類領域還是具有很大的探索價值，今后將繼續豐富該方向的實驗內容。深度學習如今發展迅猛，相信將來在垃圾分類領域會出現更好的可供選擇的分類模型。