基于改進(jìn)YOLOv5的多目標(biāo)垃圾分類研究

孫 奧，代 琦，賀平安

(浙江理工大學(xué)，浙江 杭州 310020)

0 引言

垃圾分類是“關(guān)鍵小事”和“民生大事”[1]。對于垃圾的分類，我們可以通過攝像頭采集垃圾圖像，通過垃圾分類模型快速檢測出垃圾的類別，降低人工成本，提高分類效率。因此研究垃圾分類算法有十分重要的現(xiàn)實(shí)意義。目前深度卷積神經(jīng)網(wǎng)絡(luò)在圖像分類上的精度和速度都有很大的提升。

1989 年，LeCun 提出卷積神經(jīng)網(wǎng)絡(luò)（CNN）[2]可以有效、自動地提取圖像特征，并有較快的運(yùn)行速度和識別精度。在2012 年，AlexNet 獲得了ImageNet 圖像分類競賽的冠軍，標(biāo)志著深度學(xué)習(xí)的崛起。2013 年ImageNet 圖像分類競賽的冠軍由ZFNet 獲得，并將top5錯誤率降低了近三個百分點(diǎn)。2014年的GoogLeNet在圖像識別領(lǐng)域的準(zhǔn)確度已經(jīng)達(dá)到認(rèn)類水平。2015年ResNet 模型的準(zhǔn)確度超過人類水平，并且其網(wǎng)絡(luò)層級非常深，可以提取更多的特征，標(biāo)志著深度學(xué)習(xí)進(jìn)入新的發(fā)展階段。同時隨著移動終端設(shè)備的不斷增加，也催生了輕量級的圖像識別網(wǎng)絡(luò)，如Google 團(tuán)隊(duì)于2017 年發(fā)布的MobileNet，曠世團(tuán)隊(duì)于2017 年發(fā)布ShuffleNet，Google 團(tuán)隊(duì)又于2019 年發(fā)布EfficientNet。以上輕量級框架同時兼顧了精度和識別速度，方便部署在終端設(shè)備上，使得深度學(xué)習(xí)技術(shù)應(yīng)用更為廣泛。

我國將垃圾分為可回收垃圾、廚余垃圾、有害垃圾和其他垃圾四大類。按照此分類標(biāo)準(zhǔn)的垃圾圖像識別研究還比較少，準(zhǔn)確率低，處理速度慢，并且大部分只能識別一張圖片上的單個垃圾目標(biāo)。本文通過改進(jìn)YOLOv5 多目標(biāo)檢測網(wǎng)絡(luò)，來實(shí)現(xiàn)一種更輕量、快速、準(zhǔn)確的網(wǎng)絡(luò)模型YOLOZ。

1 網(wǎng)絡(luò)模型及改進(jìn)

1.1 YOLOv5網(wǎng)絡(luò)

本文基于YOLOv5 網(wǎng)絡(luò)來實(shí)現(xiàn)多目標(biāo)的檢測，YOLOv5 的Backbone 模塊使用NewCSP-Darknet53 結(jié)構(gòu)來進(jìn)行圖像的特征提取，其中ConvBNSiLU 模塊用來對輸入圖片進(jìn)行下采樣。BottleNeck 模塊利用了Resnet網(wǎng)絡(luò)中的殘差結(jié)構(gòu)[3]，解決了隨著網(wǎng)絡(luò)深度的加深而出現(xiàn)的網(wǎng)絡(luò)退化問題。C3 模塊將原輸入分成兩個分支，分別進(jìn)行卷積操作，使得通道數(shù)減半，這樣是為了讓模型學(xué)習(xí)到更多的特征，移除計算瓶頸并降低顯存的使用[4]。SPPF 是SPP(空間金字塔池化Spatial Pyramid Pooling)[5]模塊的優(yōu)化結(jié)構(gòu)，可以加快GPU 的計算速度，并解決輸入圖像的大小必須固定的問題。Neck 模塊使用了FPN[6]加PAN[7]的結(jié)構(gòu)，F(xiàn)PN 就是一種自頂向下的側(cè)邊連接在不同尺度上構(gòu)建出的高級語義特征圖，但是這種自頂向下的FPN 網(wǎng)絡(luò)受到單項(xiàng)信息流的限制，精度有待提高。PAN 是在FPN 的基礎(chǔ)上再建立一條自下而上的通路，這樣可以讓高層的特征圖具有更強(qiáng)的語義信息，有利于物體分類，底層的特征圖具有更強(qiáng)的位置信息，有利于物體定位。Head模塊的主要部分是三個Detect 檢測器，利用基于網(wǎng)格的anchor 在不同尺度的特征圖上進(jìn)行目標(biāo)檢測，輸出目標(biāo)的坐標(biāo)信息、置信度，以及類別信息。

1.2 BatchNormalization

Normalization 就是將所有的數(shù)據(jù)映射到同一尺度，讓特征圖滿足均值為0，方差為1 的分布。理論上整個樣本集的特征圖分布都要滿足這一規(guī)律，但對于大的數(shù)據(jù)集來說這是不太現(xiàn)實(shí)的，所以采用Batch-Normalization[8]，batch 越大越接近整個數(shù)據(jù)集的分布。Batch Normalization 的計算方法如公式⑴所示，x為輸入特征圖，Mean 為特征圖的均值，Var 標(biāo)識特征圖的方差，γ(默認(rèn)值為1)用來調(diào)整數(shù)值分布的方差大小，β(默認(rèn)值為0)用來調(diào)節(jié)數(shù)值均值的位置，這兩個參數(shù)是在方向傳播過程中學(xué)習(xí)得到的。

1.3 結(jié)構(gòu)重參數(shù)化

本文提出一種結(jié)合VGG 和ResNet 的模塊RepVGGBlock[9]，在訓(xùn)練時使用一個類似ResNet-style的多分支模型，可以增加網(wǎng)絡(luò)模型的表征能力防止網(wǎng)絡(luò)出現(xiàn)退化現(xiàn)象，但是推理速度慢。而單路模型使用的3*3卷積模塊推理速度快。所以我們采用結(jié)構(gòu)重參數(shù)化的方法，在推理時把多分支模型修改為單路模型，如圖1 所示，Conv2d 為卷積模塊，BN 是BatchNormalization 模塊。對于沒有卷積模塊的分支需要構(gòu)建一個恒等的Conv2d 模塊，再把所有的Conv2d 模塊的卷積核大小調(diào)整為3，這樣就可以進(jìn)行結(jié)構(gòu)重參數(shù)化得到單路結(jié)構(gòu)。由公式⑵和公式⑶可知，Conv2d模塊和BN模塊可以融合為一個Conv2d模塊。

圖1 結(jié)構(gòu)重參數(shù)化

1.4 BiFPN

YOLOv5 模型的Neck 結(jié)構(gòu)中有些網(wǎng)絡(luò)節(jié)點(diǎn)只有一條輸入邊，如圖2左側(cè)結(jié)構(gòu)所示，刪除它對我們的網(wǎng)絡(luò)影響不大，同時可以簡化網(wǎng)絡(luò)。如果原始輸入節(jié)點(diǎn)和輸出節(jié)點(diǎn)在同一層，就在原始輸入節(jié)點(diǎn)和輸出節(jié)點(diǎn)之間添加一條額外的邊，這可以在不增加太多成本的情況下融合更多的特征[10]，可得到圖2 右側(cè)的結(jié)構(gòu)。傳統(tǒng)的特征融合往往只簡單地把不同的特征圖疊加或相加，比如Concat或Shortcut連接，而沒有對不同的特征圖進(jìn)行區(qū)分。然而輸入的各特征圖分辨率也各不同，他們對融合輸入特征圖的貢獻(xiàn)也是不同的，對其進(jìn)行簡單的相加或疊加處理并不是最佳操作，我們采用了一種類似SoftMax 的方法，把尺度歸一到[0,1]之間的簡單高效的加權(quán)特征融合機(jī)制。

圖2 BiFPN結(jié)構(gòu)

綜上我們結(jié)合結(jié)構(gòu)重參數(shù)化，BiFPN 模塊對原始YOLO 網(wǎng)絡(luò)進(jìn)行優(yōu)化，得到改進(jìn)后的網(wǎng)絡(luò)YOLOZ，如圖3所示。

圖3 YOLOZ網(wǎng)絡(luò)結(jié)構(gòu)

2 實(shí)驗(yàn)及結(jié)果分析

2.1 實(shí)驗(yàn)環(huán)境

本文實(shí)驗(yàn)在Windows10系統(tǒng)下，利用Python語言，Pytorch 深度學(xué)習(xí)框架。機(jī)器配置為CPUIntelCorei5-11600k@3.9GHz，內(nèi)存為32GB，顯卡為NvidiaGeForceRTX3050，顯存為8GB。

2.2 實(shí)驗(yàn)數(shù)據(jù)

數(shù)據(jù)來源為華為云AI 大賽提供的垃圾圖片數(shù)據(jù)以及網(wǎng)絡(luò)收集的垃圾圖片，并手動劃分四大類，分別為：可回收垃圾、有害垃圾、易腐垃圾、其他垃圾。每大類又細(xì)分為三小類。其中可回收垃圾細(xì)分為紙箱、塑料瓶、泡沫盒子，有害垃圾細(xì)分為電池、燈泡、軟膏，易腐垃圾細(xì)分為剩飯剩菜、蛋糕、蛋殼，其他垃圾細(xì)分為牙刷、口罩、煙蒂。共12 個類別3932 張圖片，8448個標(biāo)注信息。

2.3 實(shí)驗(yàn)結(jié)果

2.3.1 Mosaic數(shù)據(jù)增強(qiáng)

為了增強(qiáng)數(shù)據(jù)集的圖片和標(biāo)簽數(shù)量，可采用Mosaic數(shù)據(jù)增強(qiáng)的方式，Mosaic的方法是：使用4張圖片，對其進(jìn)行隨機(jī)裁剪和縮放旋轉(zhuǎn)等操作，最后合成1張圖像，原圖的標(biāo)簽信息也映射到新生成的圖像中，豐富了檢測物體的背景，增加了小尺度目標(biāo)的數(shù)量、數(shù)據(jù)集圖片的數(shù)量以及標(biāo)簽信息。但是Mosaic 數(shù)據(jù)增強(qiáng)是一個耗時的方法，所以我們可以設(shè)置一個閾值和一個隨機(jī)數(shù)，當(dāng)生成的隨機(jī)數(shù)小于該閾值時才采用Mosaic 數(shù)據(jù)增強(qiáng)，來根據(jù)不同的場景權(quán)衡訓(xùn)練速度和模型精度的平衡。表1 所示是設(shè)置不同閾值下的Mosaic數(shù)據(jù)增強(qiáng)對模型精度的影響。

表1 Mosaic數(shù)據(jù)增強(qiáng)結(jié)果

由表1可知，本實(shí)驗(yàn)數(shù)據(jù)較少，如果不采用Mosaic數(shù)據(jù)增強(qiáng)的方式，模型的精度很低，采用Mosaic 數(shù)據(jù)增強(qiáng)時，模型的精度相比不使用Mosaic 數(shù)據(jù)增強(qiáng)大幅提高，且隨著閾值的增大，精度不斷提高，但是模型每輪訓(xùn)練耗時也再增加。

2.3.2 結(jié)構(gòu)重參數(shù)化

本論文模型中使用了結(jié)構(gòu)重參數(shù)化技術(shù)，可以讓訓(xùn)練后的模型在推理時保持精度不變的前提下提高推理速度。表2比較了同一模型在使用該技術(shù)前后的推理速度和推理精度。其中圖片的輸入尺寸全部被調(diào)整為640*640 大小，每張圖片的推理時間是推理測試集所需時間除以測試集圖片數(shù)的結(jié)果。

表2 結(jié)構(gòu)重參數(shù)化

由表2可知，開啟結(jié)構(gòu)重參數(shù)化后，網(wǎng)絡(luò)的推理速度提升了8%，模型的計算量下降了3.2%，并且網(wǎng)絡(luò)的精度沒有下降。

2.3.3 BiFPN

本論文模型中引進(jìn)了BiFPN 網(wǎng)絡(luò)，優(yōu)化了原模型中的FPN 和PAN 結(jié)構(gòu)，可以在不增加太多成本的情況下融合更多特征。通過實(shí)驗(yàn)對比使用BiFPN 優(yōu)化前后的網(wǎng)絡(luò)，結(jié)果如表3所示。

表3 BiFPN結(jié)構(gòu)優(yōu)化網(wǎng)絡(luò)

由表3 可知，當(dāng)使用了BiFPN 時網(wǎng)絡(luò)的mAP@.5提高了1.1 個百分點(diǎn)，mAP@.5:.95 提高了2 個百分點(diǎn)，引入BiFPN 結(jié)構(gòu)相比于原始的FPN+PAN 可以提高網(wǎng)絡(luò)的精度。

2.3.4 不同多目標(biāo)檢測模型對比

基于YOLOv5 并結(jié)合本文所述的三種對網(wǎng)絡(luò)的改進(jìn)得到Y(jié)OLOZ 網(wǎng)絡(luò)，通過和前幾代YOLO 版本以及其他目標(biāo)檢測框架的比較，得出結(jié)果如表4 所示?？梢娤啾容^于YOLOv3 和YOLOv4，模型的大小顯著降低，并且精度也有所提高。相較于YOLOv5 系列，本模型體積比YOLOv5s略大，但是在精度上有很大提高，并且比YOLOv5m 體積更小，精度更高。相比較于其他目標(biāo)檢測模型如EfficientDet，在檢測速度和模型大小方面都占據(jù)優(yōu)勢。原YOLOv5 模型與YOLOZ模型在訓(xùn)練時的mAP 曲線如圖4 所示，可見YOLOZ模型在收斂速度與mAP值上均優(yōu)于原YOLOv5模型。

表4 不同目標(biāo)檢測網(wǎng)絡(luò)的對比

圖4 模型修改前后訓(xùn)練時mAP對比

3 結(jié)論

本文針對多目標(biāo)垃圾分類問題改進(jìn)了YOLOv5網(wǎng)絡(luò)的Backbone 結(jié)構(gòu)、Neck 結(jié)構(gòu)，并對數(shù)據(jù)集采用了Mosaic 的增強(qiáng)方法。得到了YOLOZ 模型，該模型在精度和體積方面都優(yōu)于改進(jìn)前的模型。但是在數(shù)據(jù)集的數(shù)量與標(biāo)注質(zhì)量上，本數(shù)據(jù)集圖片質(zhì)量較低，數(shù)量較少，且分類類別較多，導(dǎo)致模型的mAP 值不高。接下來，還需要繼續(xù)對YOLOZ 模型進(jìn)行優(yōu)化，結(jié)合ShuffleNet、EfficientNet 等輕量級網(wǎng)絡(luò)，進(jìn)一步減小模型體積，并將其移植到嵌入式設(shè)備中。