基于CBAM-EfficientNet的垃圾圖像分類算法研究

葉 沖，楊晶東

(上海理工大學 光電信息與計算機工程學院，上海200093)

0 引 言

如今，由人工進行垃圾分揀不僅對人體健康有傷害，而且垃圾分揀效率低。在日常生活中，每天會產生大量不同種類的垃圾，人工分揀只能解決其中的小部分，大多數都會進行填埋，從而對環境有很大的污染。隨著計算機視覺的發展和圖像數據的日益增多，深度學習方法在垃圾分類領域被廣泛的應用。通過對不同種類的垃圾圖像進行檢測，讓機器自動進行識別和分揀，從而提高資源利用率，減少環境污染。

傳統的圖像分類方法主要分為兩部分，一是通過人為設置特征提取器提取圖像特征，例如HOG特征［1］、SIFT特征［2］、LBP特征［3］等；二是通過設計更好的分類器算法來提高分類結果的準確率。然而隨著移動互聯網的普及以及各類互聯網產品的推出，每時每刻都有海量的圖像數據產生，這些圖像數據復雜而且多樣化，傳統的圖像分類算法并不足以支撐這些數據的分類。近年來，由于計算機技術的快速迭代和計算能力的日異豐富，給基于神經網絡學習的圖像分類算法提供了可能性。因此，越來越多的學者將目光投向基于神經網絡的圖像分類算法研究，以便找到更加快速高效的圖像分類技術。深度學習通過多層非線性層的疊加［4-5］，使得模型可以擬合更加復雜的非線性映射。2014年，來自牛津大學的Simonyan K［6］等人提出了VGG網絡，該網絡常用的有VGG16和VGG19兩種結構，除了在網絡深度上的不同，二者在本質上并無區別。相較于AlexNe［7］直接使用具有較大感受野的11*11、7*7、5*5大小的卷積核，VGG用3*3的卷積核堆疊的方法來代替。通過這種方法不僅大大減少了網絡的參數量，并且通過堆疊方式增加了網絡的非線性程度，可以使網絡擬合更加復雜的分布。同年，來自Google公司的Szegedy［8］等人提出了GoogleNet，相比于VGG網絡通過增加網絡深度來提高模型的擬合能力，GoogleNet使用了Inception模塊化結構。該結構通過不同大小的卷積核獲得不同大小的感受野，然后進行拼接，將不同尺度的特征進行融合，從而提升網絡的表達能力。2015年，Szegedy等［9］提出了Inception V2結構，該結構首先吸收了VGG網絡的優點，將多個大卷積替換為小卷積疊加，節省了大量的計算量；二是提出了批歸一化，降低了網絡對初始化權重的敏感性，并且加快了網絡的收斂速度。隨后，Szegedy等［10］對Inception結構進行了進一步的挖掘和改進，形成了Inception V3結構。該結構相比于Inception V2主要有3點改進:一是將n*n的卷積結構分解為先進行1*n卷積，然后再進行n*1卷積的結構，大大降低了網絡的參數量，從而可以堆疊更多的模塊來提升網絡的擬合能力；二是通過并行結構來優化池化部分，實現不同特征尺度的融合；三是使用標簽平滑對網絡輸出進行正則化。通過這些改進，InceptionV3結構相比于GoogleNet在ImageNet數據集上降低了約8%的錯誤率。同年，KaiMing He［11］等人提出了ResNet神經網絡，相比于VGG網絡，Resnet多了一條殘差回流通道，直接繞道將輸入信息傳給輸出，保護了信息的完整性。Resnet的殘差結構使得網絡只需學習輸入和輸出的差值，簡化了訓練目標，降低了訓練難度，一定程度上解決了梯度消失和梯度爆炸的問題。2017年，為了解決過多超參數給網絡設計和計算帶來額外難度的問題，Saining Xie等［12］提出了ResNext網絡。該網絡使用了一種平行堆疊的結構來代替ResNet的三層卷積結構，使得網絡在不明顯增加參數的同時，明顯提高了準確率。另外由于基礎拓撲結構相同，減少了大量的超參數，便于訓練和移植。為了解決單一增加網絡寬度或單一增加網絡深度導致的性能瓶頸和參數過剩問題，Tan M等［13］提出了一種搜索網絡架構Efficient-net，該網絡相比之前的卷積神經網絡，主要有Efficient-net-b0到Efficient-net-b7共8種結構分別對應于不同分辨率的圖像。

上述方法主要是基于大型數據集進行訓練，在算法上實現了較大的創新。然而現實生活中常常面對小批量的數據分類任務，由于網絡參數過多、太少的數據集，使以上網絡模型難以得到充分的訓練或者容易出現過擬合現象。雖然遷移學習可以解決這一問題，但面對不同的數據集任務，容易忽略數據之間的差異性。

因此，本文提出了一種自適應注意力機制和數據增強方法。通過數據增強，提高數據集的多樣性，降低過擬合效應；在原模型中添加自適應注意力機制，通過自適應注意力機制，提取目標數據集較重要的特征，最終獲得準確率的提高。

本文的主要貢獻如下:

（1）提出一種數據增強的方法，幫助擴展訓練集的多樣性；

（2）將EfficientNet、VGG和ResNet進行對比，說明了EfficientNet的高效性和將其作為垃圾分類遷移學習主干網絡的正確性；

（3）將CBAM注意力機制用于EfficientNet遷移學習網絡中，并通過Grad-CAM［14］對原圖信息的重要性進行可視化，說明了本文所提模型能夠加強特征提取功能，從而得到較優的分類準確率。

1 模型架構

1.1 Efficient-net網絡模型

經典的神經網絡一般有以下3個特點:一是利用殘差神經網絡來增加網絡的深度，通過更深的神經網絡層數來提取更深的特征，并獲取一定的性能提升；二是通過改變每一層提取的特征層數，實現更多的特征提取，得到更多的特征，以此來增加網絡的多尺度表達特性；三是通過增大輸入圖像的分辨率，來幫助網絡學到更多的圖像細節。Efficient-net是將這3個特點結合起來，通過調整輸入圖像的分辨率、深度、寬度3個維度，來得到更好的網絡結構。其基線結構如圖1所示。

圖1 EfficienNet B0結構圖Fig.1 Structure of EfficienNet B0

在該網絡的基礎上，作者通過復合縮放的方法對網絡的深度、寬度以及輸入圖像的分辨率3個維度進行優化，以獲得在一定資源條件限制下的準確率最高的模型。其對縮放的關系如式（1）所示:

其中，a、β、γ是通過神經網絡結構搜索得到的常數；?是根據計算資源大小設置的常數；d、w、r分別代表網絡的深度、寬度和分辨率的縮放系數。Efficient-net-b1～b7是通過確定a、β、γ的最優取值后，調整?所得到的。因此Efficient-net在有限的資源環境下，可以獲得良好的性能提高。

1.2 CBAM注意力機制

CBAM注意力機制［15］是一種簡單而高效的注意力模塊。其將給定的中間特征圖沿空間和通道2個獨立的維度依次判斷特征注意力圖，并且與原始特征圖相乘進行自適應優化，其結構如圖2所示。

圖2 CBAM模型結構圖Fig.2 Structure of CBAM

1.3 整體模型結構

由于目標數據集與源數據集之間的差異，僅通過遷移學習提取特征容易造成特征提取不充分，從而導致最終精度的損失。一般做法是，放開所有預訓練網絡參數進行微調，但由于目標數據集數量的不足，該方法較容易得到過擬合的模型，不具有很強的泛化性能。因此，本文在預訓練網絡中的每一次下采樣前加入CBAM注意力機制，以此來提升網絡對于某些重要特征圖和重要空間的注意力，從而提高模型準確率。結構如圖3所示。

圖3 基于CBAM注意力機制的EfficientNet結構圖Fig.3 Eficientnet structure based on CBAM attention mechanism

1.4 損失函數

由于垃圾圖像數據集中，存在類別不平衡以及少部分圖像未正確標注的問題，本文采用了標簽平滑（Label Smooth）［16］和Focal loss［17］相結合的損失函數。標簽平滑是假定標簽并不是100%正確，將預定的類別設置為一個較大的概率，其它類別分配相應較小的概率。在存在較多分類類別、標注異常時，會起到較大的改善作用。Focal loss是通過動態增加難分類樣本的損失函數權重，降低易分類樣本的損失函數權重，使得模型著重于難訓練樣本的訓練，從而緩解樣本不平衡的問題。如式（2）所示:

1.5 訓練策略

首先，將EfficientNet網絡去除最后一層全連接層，然后使用Google公司在ImageNet上訓練完成的EfficientNet B0權重將其初始化；其次，加入自適應注意力機制重構網絡，并將其余權重采用He＿norma的方式進行初始化；固定前3層網絡，采用本文所提的損失函數進行訓練；最后在驗證集loss最低處保存模型。具體訓練步驟如下:

輸入:形狀為［N，H，W，C］的圖像和標簽，

Step1:構建EfficientNet網絡；

Step2:加載ImageNet預訓練參數；

Step3:加入自適應注意力機制重構網絡；

Step4:固定網絡前3層進行訓練；

Step5:在驗證集loss最低處保存網絡模型。

輸出:圖像的分類結果。

2 實驗結果分析

本實驗平臺為兩塊GPU顯卡，GPU型號為GeForce RTX 2080 Ti，該GPU顯存為11.07GB，顯卡頻率1.112（GHZ）。實驗基于深度學習框架為keras2.2.4和tensorflow1.14，使用python3.5語言編程。

2.1 數據集描述

本文所采用的數據集來自于2019年華為杯垃圾分類挑戰賽，該數據集共有14 802張圖片，分為40個類別。圖4為該數據集的統計結果。由圖中可以發現，該數據集主要有以下特點:一是各類樣本數量不均衡，其中第4類樣本最少，僅有100張圖片不到；二是圖片分辨率大小不一致，包含多種分辨率的圖像。由于本文采取Efficient-net b0提取圖像特征，因此本文經過預處理將圖片統一為300×300分辨率大小的圖片。并將數據集劃分為訓練集、驗證集和測試集，其劃分比例為7∶2∶1。

圖4 數據集統計圖片Fig.4 Statistical picture of data set

2.2 數據集增強

本文首先按照圖像最大邊長與輸入圖像分辨率（300），進行等比拉伸，然后將圖像進行隨機裁剪、顏色扭曲、圖像旋轉等進行圖像增強。最后，將圖像邊界用0填充至輸入圖像分辨率大小（300*300），如圖5所示。左邊為數據集部分原圖，中間為本文所提數據增強方法增強后結果圖，右邊為直接resize后的結果圖。從圖中可以發現，經本文所提方法增強后，其圖像比例與原始圖像保持一致，而直接resize后其圖像有所失真，不利于模型識別。

圖5 數據增強結果圖Fig.5 Data enhancement results

2.3 實驗結果與分析

利用Efficient-net網絡與VGG和Resnet在該數據集上進行了實驗。為驗證本文所采用遷移學習網絡在垃圾分類數據集上的高效性和準確性，本文將訓練階段驗證集準確率和驗證集loss的變化進行了可視化，如圖6、圖7所示。由圖可知，Efficientnet采用了更高效網絡特征提取模塊，對未來實現垃圾分類自動化分揀有一定現實意義。其網絡相比于VGG和Resnet收斂較快，并且有較高的準確率。在驗證集上相比于Resnet提高了5.1%，相比于VGG其準確率提高了5.8%，說明了Efficient-net網絡在垃圾分類數據集上的有效性和高效性。

圖6 網絡訓練正確率對比圖Fig.6 Comparison of network training accuracy

圖7 網絡訓練loss對比圖Fig.7 Comparison of network training loss

為進一步說明自適應注意力機制的有效性，本文采用了Grad-CAM［14］（Grad-CAM是最近提出的一種可視化方法，其使用梯度來計算卷積層中空間位置的重要性）對結果圖進行可視化，如圖8所示。其中，圖8（a）為原圖，圖8（b）為EfficientNet最后一層特征層可視化結果，圖8（c）為在預訓練網絡中加入自適應注意力后，最后一層特征可視化結果。由圖8可以發現，加入注意力機制后，其特征提取更為精確，且該圖像所屬類別的置信度更高，進一步說明了基于自適應注意力機制在EfficientNet網絡的有效性。

圖8 圖像空間重要性可視化結果圖Fig.8 Visualization results of image spatial importance

最后，本文對比了Efficient-net和CBAMEfficientNet在測試集上的準確率和模型大小，結果，見表1。由表1可以發現，加入CBAM注意力機制后，雖然模型參數量有所增加，但是計算速度卻降低了近1%，并且準確率提升了2.5%，證明了CBAMEfficientNet在垃圾圖像分類數據集上的有效性。

表1 模型結果對比Tab.1 Comparison of model results

3 結束語

本文將EfficientNet網絡應用于垃圾圖像分類數據集中，并與VGG和Resnet網絡進行了對比。實驗結果顯示，EfficientNet相較于VGG和Resnet網絡具有較高的準確率和高效性，適合輕量化部署。針對遷移學習特征提取不充分問題，本文在預訓練網絡中加入了CBAM注意力機制，增強了預訓練網絡中重要的特征層權重，同時抑制無效特征層的影響，并與EfficientNet進行對比，實驗結果證明CBAMEfficientNet相比于EfficientNet提高了2.5%的準確率，且計算速度未明顯降低。進一步的工作將集中于更加適合遷移學習的注意力機制模型研究，目的是進一步增強網絡的特征提取能力，從而得到較好的性能指標。