面向垃圾分類場(chǎng)景的輕量化目標(biāo)檢測(cè)方案

陳健松，蔡藝軍

(廈門理工學(xué)院 光電與通信工程學(xué)院，福建 廈門 361024)

隨著人民生活水平的提高，生活垃圾數(shù)量迅猛劇增，大量的垃圾對(duì)人類的健康、環(huán)境的破壞及經(jīng)濟(jì)的發(fā)展帶來(lái)嚴(yán)重的威脅.傳統(tǒng)的垃圾回收需要大量的人力和物力，利用深度學(xué)習(xí)技術(shù)，對(duì)垃圾進(jìn)行智能化、自動(dòng)化處理將成為一種趨勢(shì)，可以提高垃圾的處理效率，有利于綠色經(jīng)濟(jì)和低碳循環(huán)發(fā)展.

近年來(lái)，基于深度學(xué)習(xí)的目標(biāo)檢測(cè)算法已經(jīng)成為主流，深度學(xué)習(xí)以其強(qiáng)大的特征表示能力和端到端的學(xué)習(xí)能力，極大地提高了目標(biāo)檢測(cè)的精度和效率.其中目標(biāo)檢測(cè)算法主要分為一階段和二階段算法兩大類.二階段算法在第一階段專注于找到物體的位置，得到建議框，提高準(zhǔn)確率和召回率；第二階段專注于對(duì)建議框進(jìn)行物體分類，找到更加精確的錨框，其算法有R-CNN、SPPNet、Fast R-CNN 和 Faster R-CNN[1-4].一階段算法減少了生成建議框這個(gè)階段，直接產(chǎn)生物體的類別概率和位置坐標(biāo)值，經(jīng)過(guò)單次檢測(cè)可以直接得到最終的檢測(cè)結(jié)果，其代表算法有Yolo[5]、SSD[6].這2 種算法的處理過(guò)程不同，在性能表現(xiàn)上的差異較大.前者在檢測(cè)精度和定位精度方面具有優(yōu)勢(shì)，后者在速度方面具有優(yōu)勢(shì).

Yolo 系列算法作為一階段算法，因其優(yōu)秀的檢測(cè)性能得到了廣泛的應(yīng)用，對(duì)于邊緣端設(shè)備十分友好.研究人員對(duì)Yolo 系列算法提出許多改進(jìn)方法.李小波等[7-8]在模型中添加注意力機(jī)制，增強(qiáng)特征圖的信息表達(dá)能力，提高模型對(duì)小尺度目標(biāo)的檢測(cè)能力，但是模型的復(fù)雜度相應(yīng)增加，不利于邊緣端部署.李仁鷹等[9-11]通過(guò)將模型的主干網(wǎng)絡(luò)替換為MobileNetV2，并嵌入注意力機(jī)制來(lái)實(shí)現(xiàn)模型的輕量化，在模型復(fù)雜度上有了改進(jìn)，但沒有很好地兼顧精度.王相友等[12-13]利用通道剪枝來(lái)實(shí)現(xiàn)模型的輕量化，楊小岡等[14]在將模型的主干網(wǎng)絡(luò)修改成MobileNetV3 的基礎(chǔ)上進(jìn)行迭代剪枝.這2 種改進(jìn)方式都使模型的推理速度大大加快，但是精度有所下降.

以上研究只針對(duì)精度和輕量化其中一種進(jìn)行改進(jìn)，不能同時(shí)滿足檢測(cè)精度和推理速度2 種要求.本文對(duì)Yolov5s 模型進(jìn)行輕量化改進(jìn)，在保證模型精度的同時(shí)，提高模型的推理速度，對(duì)于邊緣端設(shè)備更加友好.

1 Yolov5 算法改進(jìn)

從模塊和數(shù)據(jù)2 個(gè)方面進(jìn)行改進(jìn).如圖1 所示，在模塊方面，引入Stem 模塊和深度可分離卷積，對(duì)C3 模型進(jìn)行改進(jìn).在數(shù)據(jù)方面，利用Kmeans++算法重新計(jì)算物體的錨框值.

圖1 Yolov5 算法的改進(jìn)樹狀圖Fig.1 Improvement tree graph of Yolov5 algorithm

1.1 改進(jìn)模型的結(jié)構(gòu)圖

以Yolov5 結(jié)構(gòu)為基礎(chǔ)，對(duì)模型的主干網(wǎng)絡(luò)和頸部網(wǎng)絡(luò)進(jìn)行修改.模型結(jié)構(gòu)如圖2 所示，模型主要由輸入(Input)、主干網(wǎng)絡(luò)(Backbone)、頸部網(wǎng)絡(luò)(Neck)和檢測(cè)頭(Detect)4 部分組成.將Stem 模塊作為輸入圖像處理模塊，能夠在幾乎不增加模型參數(shù)量的同時(shí)，提高模型的特征提取能力，因此采用Stem 模塊作為圖像輸入的處理模型[15].為了減少模型的復(fù)雜度，提高模型的推理速度，將主干網(wǎng)絡(luò)和頸部網(wǎng)絡(luò)中的3×3 降采樣普通卷積改成深度可分離卷積.為了進(jìn)一步減少參數(shù)量和模型復(fù)雜度，將頸部結(jié)構(gòu)中的C3 模塊改進(jìn)成C3-C 模塊來(lái)減少參數(shù)量，將backbone 中的C3 模塊改進(jìn)成C3-B 模塊來(lái)增強(qiáng)圖像特征的提取能力.

圖2 改進(jìn)后的Yolov5 模型結(jié)構(gòu)圖Fig.2 Structure diagram of improved Yolov5 model

1.1.1 Stem 模塊 Stem 模塊來(lái)源于移動(dòng)端的實(shí)時(shí)目標(biāo)檢測(cè)模型Peleenet[16].圖3 給出Stem 模塊結(jié)構(gòu)圖.可以看出，Stem 模塊首先進(jìn)行快速降維操作，采用卷積核為3×3、步長(zhǎng)(stride)為2 的卷積層實(shí)現(xiàn).采用兩分支的結(jié)構(gòu)：一個(gè)分支用步長(zhǎng)為1、卷積核為1×1 的卷積層和步長(zhǎng)為2、卷積核為3×3的卷積層，另一個(gè)分支使用步長(zhǎng)為2、卷積核為2×2 的最大池化層(Maxpool).最大池化層在保留主要特征的同時(shí)，減少了參數(shù)量和計(jì)算量，防止了過(guò)擬合的發(fā)生.Stem 模塊在進(jìn)行第一次降采樣后，分成2 條路徑對(duì)圖像特征進(jìn)行提取，使得模型可以提取更加豐富的特征.這樣可以在不帶來(lái)過(guò)多計(jì)算耗時(shí)的前提下，提高網(wǎng)絡(luò)的特征表達(dá)能力.

圖3 Stem 模塊的結(jié)構(gòu)圖Fig.3 Structural diagram of Stem module

1.1.2 深度可分離卷積 一些輕量級(jí)的網(wǎng)絡(luò)，如MobileNet 網(wǎng)絡(luò)中，使用深度可分離卷積(depthwise separable convolution)來(lái)提取圖像特征[17-19].在特征提取的過(guò)程中，隨著網(wǎng)絡(luò)深度的增加，特征圖的通道數(shù)不斷增加，當(dāng)使用傳統(tǒng)卷積對(duì)深層特征圖提取特征時(shí)，必然會(huì)產(chǎn)生較大的參數(shù)量，導(dǎo)致算法的計(jì)算運(yùn)行速度較慢.與常規(guī)的卷積操作相比，深度可分離卷積的參數(shù)量和運(yùn)算成本較小.

深度可分離卷積主要分為2 個(gè)過(guò)程，分別為逐通道卷積(depthwise convolution) 和逐點(diǎn)卷積(pointwise convolution).如圖4(a)所示為深度可分離卷積的實(shí)現(xiàn)，深度可分離卷積首先進(jìn)行逐通道卷積，其卷積核數(shù)量與輸入的通道數(shù)相同，每個(gè)卷積核負(fù)責(zé)一個(gè)通道，經(jīng)過(guò)卷積后得到的特征圖的通道數(shù)與輸入的通道數(shù)一致.由于逐通道卷積沒有利用特征圖的空間特征，需要利用逐點(diǎn)卷積，逐點(diǎn)卷積相當(dāng)于1×1 的普通卷積(普通卷積如圖4(b)所示)，它的卷積核尺寸為1×1×M(M為輸入的通道數(shù))，這里的卷積運(yùn)算會(huì)將上一步的特征在深度方向上進(jìn)行加權(quán)組合，生成新的特征圖.

圖4 深度可分離卷積模塊的結(jié)構(gòu)圖Fig.4 Structural diagram of depthwise separable convolution module

傳統(tǒng)卷積的計(jì)算量為

深度可分離卷積的計(jì)算量為

式中：DK×DK為卷積核的尺寸，M為卷積核的深度，DW×DH為輸出特征圖的尺寸，N為輸出特征圖的深度.深度可分離卷積分為深度卷積和逐點(diǎn)卷積2 個(gè)階段，相比于傳統(tǒng)卷積，深度可分離卷積在完成相同計(jì)算目的的情況下計(jì)算開銷更小.深度可分離卷積與傳統(tǒng)卷積的計(jì)算量之比為

以3×3 卷積為例，深度可分離卷積的參數(shù)量可以減少至普通卷積的(1/9+1/N)倍，隨著網(wǎng)絡(luò)層數(shù)的增加，通常會(huì)導(dǎo)致模型參數(shù)量的顯著增加，但使用深度可分離卷積能夠有效地減少模型參數(shù)量.

1.1.3 C3 模塊改進(jìn) 對(duì)于在嵌入式端實(shí)現(xiàn)的垃圾分類識(shí)別，既要應(yīng)對(duì)算力不足的情況，還要考慮復(fù)雜環(huán)境下對(duì)物體的識(shí)別準(zhǔn)確率，因此對(duì)主干網(wǎng)絡(luò)和頸部網(wǎng)絡(luò)的C3 模塊進(jìn)行改進(jìn).如圖5(a)所示為主干網(wǎng)絡(luò)使用的bottleneck 模塊.該模塊包含了跨界融合層，又包含了多個(gè)Conv2d 卷積層，因而該模塊既可以實(shí)現(xiàn)對(duì)圖像低層特征和高層特征的融合(可以彌補(bǔ)低分辨率的高層特征圖中空間信息的損失)，又能夠?qū)崿F(xiàn)對(duì)圖像更深層特征的提取.在主干網(wǎng)絡(luò)的C3 模塊中使用了2 個(gè)Bottle-Neck-1 模塊(見圖5(b))，使得改進(jìn)后C3 模塊的輸出特征圖含有更強(qiáng)的語(yǔ)義信息與更多的細(xì)節(jié)信息，以提升該模塊的特征提取能力[20].如圖5(c)所示，該模塊將頸部網(wǎng)絡(luò)中bottleneck 模塊的普通3×3 卷積改成深度可分離卷積(DWConv)，將修改后的bottleneck 模塊應(yīng)用于頸部網(wǎng)絡(luò)的C3 模塊中，形成C3-C 模塊(見圖5(d))，減少了模型的參數(shù)量.

圖5 C3 改進(jìn)模塊的結(jié)構(gòu)圖Fig.5 Structural diagram of C3 improved module

1.2 數(shù)據(jù)預(yù)處理

Yolov5 使用 K-means 算法對(duì)通用目標(biāo)檢測(cè)數(shù)據(jù)集COCO 進(jìn)行聚類，獲得初始的先驗(yàn)錨框參數(shù).利用K-means 算法，隨機(jī)分配初始聚類中心，可能會(huì)導(dǎo)致初始聚類中心與最佳聚類中心相差甚遠(yuǎn)，且初始的先驗(yàn)錨框參數(shù)僅適用于COCO、VOC 和ImageNet 等開放場(chǎng)景數(shù)據(jù)集，而本文數(shù)據(jù)集的檢測(cè)對(duì)象與其存在較大的差異[21].選擇 K-means++算法對(duì)錨框進(jìn)行重新聚類，生成更優(yōu)的錨框值，不會(huì)帶來(lái)額外的參數(shù)和計(jì)算量.

利用K-means++算法，隨機(jī)選擇一個(gè)樣本，作為當(dāng)前的初始聚類中心；計(jì)算每個(gè)樣本和已有聚類中心之間的最短距離，將該樣本歸類到與其相距最小的聚類中心所對(duì)應(yīng)的類別中[22-23].計(jì)算每個(gè)樣本被設(shè)定為下一個(gè)聚類中心的概率，選擇概率最大的樣本作為下一個(gè)中心.概率的計(jì)算公式為

式中：D(x)為每一個(gè)樣本點(diǎn)到當(dāng)前中心的最短距離，每次分配一個(gè)對(duì)象，會(huì)根據(jù)現(xiàn)有的聚類對(duì)象重新計(jì)算聚類中心，重復(fù)該過(guò)程直到?jīng)]有對(duì)象被重新分配到其他的聚類中，最后篩選出K 聚類中心.盡管K-means++算法中用于初始選擇聚類中心的方法需要花費(fèi)額外的時(shí)間，但是K-means 算法在選取合適的聚類中心后會(huì)很快收斂，使用Kmeans++算法選擇物體的錨框值，使得Yolov5 算法可以選擇最具代表性的錨框值，使其在訓(xùn)練過(guò)程中能夠更好地預(yù)測(cè)目標(biāo)框的大小，提高模型的準(zhǔn)確性和穩(wěn)定性.

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

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

使用自己拍攝的垃圾分類數(shù)據(jù)集和從網(wǎng)上爬取的數(shù)據(jù)集，重新組合成包含3 305 張圖片的垃圾分類數(shù)據(jù)集.該數(shù)據(jù)集涵蓋了6 種常見垃圾（瓶子、電池、筆、口罩、易拉罐和紙杯），每張圖片的物品標(biāo)注數(shù)量不同，這樣的數(shù)據(jù)集可以更好地反映現(xiàn)實(shí)場(chǎng)景.為了獲得更準(zhǔn)確的訓(xùn)練效果，將數(shù)據(jù)集按照9∶1 的比例進(jìn)行劃分，其中90% 的數(shù)據(jù)，即2 975 張圖片用于訓(xùn)練，剩下的10% 的數(shù)據(jù)，即330 張圖片用于驗(yàn)證.每個(gè)類別的訓(xùn)練樣本數(shù)量如下：瓶子1 571 個(gè)、電池1 850 個(gè)、筆812支、口罩902 個(gè)、易拉罐946 個(gè)、紙杯887 個(gè).

2.2 實(shí)驗(yàn)參數(shù)配置

模型的訓(xùn)練采用pytorch1.10.1 深度學(xué)習(xí)框架，CUDA 10.1.硬件配置為GeForce RTX 3080 顯卡、10 GB 顯存.訓(xùn)練時(shí)的參數(shù)配置如表1 所示.

表1 超參數(shù)設(shè)置

2.3 實(shí)驗(yàn)評(píng)價(jià)指標(biāo)

在上述統(tǒng)一的實(shí)驗(yàn)配置環(huán)境中進(jìn)行對(duì)比實(shí)驗(yàn)，采用召回率R、精確率P、均值平均精度(mean average precision,mAP)、模型參數(shù)量和模型推理速度作為評(píng)價(jià)指標(biāo).mAP@0.5 為IOU 設(shè)為0.5 時(shí)，計(jì)算每一類的所有圖片的AP，對(duì)所有類別的AP 求平均值.mAP@0.5：0.95 表示在不同IOU 閾值(0.5～0.95，步長(zhǎng)為0.05)下的平均mAP[24].mAP 的計(jì)算如下所示：

式中：TP為正確識(shí)別出的垃圾樣本數(shù)量，F(xiàn)N為模型未能識(shí)別出的垃圾樣本數(shù)量，F(xiàn)P為模型錯(cuò)誤識(shí)別的垃圾樣本數(shù)量；P為精確率，表示模型預(yù)測(cè)為垃圾的樣本中預(yù)測(cè)正確的占比；R為召回率，表示所有的垃圾樣本中模型能夠識(shí)別出來(lái)的數(shù)量占比；n為垃圾樣本的類別數(shù)；i為當(dāng)前類別的序號(hào).

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

如圖6 所示為Yolov5s 模型和改進(jìn)模型在使用K-means++算法后精度和錨框損失L的變化情況.圖中，N為迭代次數(shù).圖中，K-means++表示使用該算法調(diào)整數(shù)據(jù)集中的錨框值進(jìn)行訓(xùn)練.從mAP的表現(xiàn)（見圖6(a)、(b)）來(lái)看，在加入K-means++算法后，模型的精度收斂速度更快，mAP 比未使用K-means++算法的模型高.如圖6(c)所示為模型的錨框損失情況.可以看出，在加入K-means++算法后，損失的收斂速度更快；與未使用K-means++算法的模型相比，錨框損失更小.

圖6 Yolov5s 模型和改進(jìn)模型的性能對(duì)比圖Fig.6 Performance comparison between Yolov5s model and improved model

如表2 所示，改進(jìn)模型和Yolov5s 模型在加入該算法后，準(zhǔn)確率、召回率和平均精度都明顯提升.具體來(lái)說(shuō)，Yolov5s 模型在加入K-means++算法后，與未加入該算法相比，準(zhǔn)確率提升了2.1%，召回率提升了4.4%，mAP@0.5 提升了2.2%，mAP@0.5:0.95 提升了2.6%.改進(jìn)模型在加入K-means++算法后，與未加入該算法相比，準(zhǔn)確率提升了3.8%，召回率提升了3.4%，mAP@0.5 提升了2%，mAP@0.5:0.95提升了2.7%.這說(shuō)明K-means 算法對(duì)錨框進(jìn)行優(yōu)化的有效性，能夠顯著提高模型的準(zhǔn)確率和穩(wěn)定性.改進(jìn)模型加上K-means++算法后，與未加入Kmeans++算法的Yolov5s 模型相比，檢測(cè)任務(wù)的表現(xiàn)更加優(yōu)越.具體來(lái)說(shuō)，加上K-means++算法的改進(jìn)模型與未加入K-means++算法的Yolov5s 模型相比，準(zhǔn)確率提升了6.3%，召回率提升了1.9%，mAP@0.5 提升了0.8%，mAP@0.5:0.95 提升了3%.這表明最終改進(jìn)的模型能夠更好地適應(yīng)各種場(chǎng)景和數(shù)據(jù)集，表現(xiàn)更加優(yōu)異.改進(jìn)模型在參數(shù)量上比Yolov5s 模型小，參數(shù)量減少至原模型的77.1%.將改進(jìn)后的模型部署到Jetson Nano 上，每幀的推理時(shí)間達(dá)到106.1 ms，與Yolov5s 原模型相比，速度提高了21.9%.這表明改進(jìn)后的模型具有更高的計(jì)算效率和更好的實(shí)用性，能夠更好地滿足實(shí)際需求.

表2 Yolov5s 模型和改進(jìn)模型的實(shí)驗(yàn)結(jié)果對(duì)比Tab.2 Comparison of experimental results between Yolov5s model and improved model

將使用了K-means++算法的改進(jìn)模型稱為最終改進(jìn)模型，未添加K-means++算法的Yolov5s 模型稱為Yolov5s 原模型.表2 中，Np為模型的參數(shù)量，t為模型的推理時(shí)間.從表2 的實(shí)驗(yàn)結(jié)果可以看出，最終改進(jìn)模型的性能優(yōu)于Yolov5s 原模型.

為了進(jìn)一步對(duì)比最終改進(jìn)模型和Yolov5s 原模型的性能，從網(wǎng)上隨機(jī)挑選圖片進(jìn)行測(cè)試.如圖7(a)所示為Yolov5s 原模型的檢測(cè)效果，如圖7(b)所示為最終改進(jìn)模型的檢測(cè)效果.可以看出，最終改進(jìn)模型對(duì)物體的識(shí)別精度有所提高，減少了誤檢率，提高了對(duì)遮擋物體的識(shí)別精度.實(shí)驗(yàn)結(jié)果表明，最終改進(jìn)模型在垃圾分類檢測(cè)任務(wù)上的表現(xiàn)更優(yōu)秀.

圖7 Yolov5s 模型和改進(jìn)模型的檢測(cè)結(jié)果對(duì)比圖Fig.7 Comparison of detection results between Yolov5s and improved model

3 結(jié)語(yǔ)

針對(duì)邊緣端算力不足導(dǎo)致目標(biāo)檢測(cè)模型實(shí)時(shí)推理性能差的問(wèn)題，本文提出輕量化的Yolov5 模型.將自制的私有數(shù)據(jù)集和網(wǎng)上爬取的數(shù)據(jù)重新組合成垃圾分類數(shù)據(jù)集.在數(shù)據(jù)上，使用K-means++算法優(yōu)化錨框；在模型上，引入Stem 模塊提高特征提取能力，利用深度可分離卷積減少模型的參數(shù)量和復(fù)雜度.對(duì)主干網(wǎng)絡(luò)的C3 模型進(jìn)行改進(jìn)，加強(qiáng)了特征提取能力.實(shí)驗(yàn)表明，改進(jìn)模型與Yolov5s 模型相比，在提高模型精度和穩(wěn)定性的同時(shí)，明顯提升了推理速度，減少了攝像頭實(shí)時(shí)檢測(cè)卡頓的情況.本文的研究為邊緣端的垃圾分類系統(tǒng)開發(fā)提供了科學(xué)設(shè)計(jì)的依據(jù)，具有廣闊的應(yīng)用前景.