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

葉 沖，楊晶東

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

0 引 言

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

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

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

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

本文的主要貢獻如下:

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

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

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

1 模型架構(gòu)

1.1 Efficient-net網(wǎng)絡模型

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

圖1 EfficienNet B0結(jié)構(gòu)圖Fig.1 Structure of EfficienNet B0

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

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

1.2 CBAM注意力機制

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

圖2 CBAM模型結(jié)構(gòu)圖Fig.2 Structure of CBAM

1.3 整體模型結(jié)構(gòu)

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

圖3 基于CBAM注意力機制的EfficientNet結(jié)構(gòu)圖Fig.3 Eficientnet structure based on CBAM attention mechanism

1.4 損失函數(shù)

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

1.5 訓練策略

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

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

Step1:構(gòu)建EfficientNet網(wǎng)絡；

Step2:加載ImageNet預訓練參數(shù)；

Step3:加入自適應注意力機制重構(gòu)網(wǎng)絡；

Step4:固定網(wǎng)絡前3層進行訓練；

Step5:在驗證集loss最低處保存網(wǎng)絡模型。

輸出:圖像的分類結(jié)果。

2 實驗結(jié)果分析

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

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

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

圖4 數(shù)據(jù)集統(tǒng)計圖片F(xiàn)ig.4 Statistical picture of data set

2.2 數(shù)據(jù)集增強

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

圖5 數(shù)據(jù)增強結(jié)果圖Fig.5 Data enhancement results

2.3 實驗結(jié)果與分析

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

圖6 網(wǎng)絡訓練正確率對比圖Fig.6 Comparison of network training accuracy

圖7 網(wǎng)絡訓練loss對比圖Fig.7 Comparison of network training loss

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

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

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

表1 模型結(jié)果對比Tab.1 Comparison of model results

3 結(jié)束語

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