基于模型剪枝和知識蒸餾的船舶目標(biāo)檢測方法

邱 春,路 紅,秦彬鑫,萬文明,戴云周

(南京工程學(xué)院機(jī)械工程學(xué)院, 江蘇 南京 211167)

近年來,隨著目標(biāo)檢測的發(fā)展,對湖面和海面的船舶目標(biāo)檢測技術(shù)研究越來越受到重視,其在軍事與民用方面都具有重要意義[1].目前,基于深度學(xué)習(xí)的目標(biāo)檢測技術(shù)需要大量的計算,卷積網(wǎng)絡(luò)里大量的參數(shù)占用內(nèi)存空間.為了解決神經(jīng)網(wǎng)絡(luò)中計算量大、占用內(nèi)存多等問題,本文采用剪枝技術(shù)來對模型進(jìn)行壓縮[2],并用知識蒸餾技術(shù)對剪枝后的模型進(jìn)行蒸餾,補(bǔ)償剪枝后模型的檢測精度損失.

1 模型選取

1.1 卷積神經(jīng)網(wǎng)絡(luò)

卷積神經(jīng)網(wǎng)絡(luò)常用于分類與檢測等任務(wù)[3],主要由卷積層、池化層、全連接層組成.卷積層是通過感受野即卷積核在圖像上進(jìn)行滑動卷積,以此來提取圖像中不同維度的特征;將卷積層輸出的特征放到池化層中對特征圖進(jìn)行池化操作,池化層與卷積層相似,是在感受野上建立一個窗口來進(jìn)行滑動操作,從而達(dá)到特征降維的作用;將池化后輸出的特征送到全連接層中進(jìn)行維度的轉(zhuǎn)換,在全連接層中對特征進(jìn)行分類等操作,輸出不同類型的特征.典型的卷積神經(jīng)網(wǎng)絡(luò)有VGG、AlexNet、ResNet等.

1.2 YOLO v3算法介紹

在計算機(jī)視覺領(lǐng)域中,基于深度學(xué)習(xí)的目標(biāo)檢測算法主要分為兩階段算法和單階段算法:1) 兩階段算法有R-CNN、Faster R-CNN算法,先獲得候選區(qū)域(region proposal),然后將其分成不同的目標(biāo)類別,該類算法的定位和檢測準(zhǔn)確率較高,但運算速度較慢,實時性達(dá)不到要求;2) 單階段算法有YOLO、SSD、YOLO900算法等,依據(jù)回歸思想使實時檢測成為可能,但檢測精度比兩階段算法低.本研究模型選擇單階段算法中的YOLO v3算法.YOLO v3算法主要由特征提取網(wǎng)絡(luò)Darknet53和檢測網(wǎng)絡(luò)組成,Darknet53網(wǎng)絡(luò)由53個卷積層組成,檢測網(wǎng)絡(luò)是在不同尺度的特征圖上選取3個尺度的特征圖進(jìn)行目標(biāo)預(yù)測輸出[4].YOLO v3算法檢測原理是:首先將輸入的圖像送到Darknet53網(wǎng)絡(luò)中進(jìn)行特征提取,得到多個尺度不一的特征圖,對特征圖進(jìn)行處理獲得目標(biāo)的相關(guān)類別與回歸框;然后經(jīng)過極大值抑制與類別置信度對目標(biāo)邊框進(jìn)行過濾;最后得到合格的目標(biāo)信息.YOLO v3算法檢測精度高、速度快,滿足了人們對目標(biāo)檢測的實時性要求.YOLO v3算法的網(wǎng)絡(luò)結(jié)構(gòu)圖如圖1所示.

圖1 YOLO v3算法網(wǎng)絡(luò)結(jié)構(gòu)圖

2 船舶檢測模型設(shè)計

2.1 模型剪枝

基于深度學(xué)習(xí)的模型極其復(fù)雜,實際應(yīng)用中往往會受到計算資源消耗大、存儲空間要求高等阻礙,對于嵌入到微型機(jī)或者移動端來說更加困難.神經(jīng)網(wǎng)絡(luò)模型中有大量的神經(jīng)元與權(quán)重冗余,真正對模型產(chǎn)生影響的只有5%～10%.將模型中不重要的參數(shù)去掉來減小模型大小的方法越來越多,最常見的模型壓縮方法是模型剪枝[5-6].剪枝不僅可以降低模型復(fù)雜度、防止過擬合,還可以壓縮模型大小、提高模型的速度.剪枝方式可分為通道剪枝、權(quán)值剪枝、層剪枝.層剪枝效果最粗糙,但無軟硬件要求、靈活性好;權(quán)值剪枝靈活性最好,但是對軟硬件有要求,需要專門的系數(shù)矩陣運算庫或者硬件來做加速;通道剪枝集成了層剪枝和權(quán)值剪枝的優(yōu)點,精細(xì)且靈活性好.本研究采用通道與層相結(jié)合方式來對模型進(jìn)行剪枝.模型剪枝示意圖如圖2所示.

圖2 模型剪枝示意圖

通道剪枝對軟硬件要求不高,在細(xì)粒度方面進(jìn)行了很好的平衡,采用過隨機(jī)拋棄channel-wise的連接引入稀疏化,將常見的L1約束添加到各個通道的縮放因子上,使其能夠更加便捷地優(yōu)化目標(biāo),通道剪枝的目標(biāo)函數(shù)為:

(1)

式中:W為訓(xùn)練的權(quán)值;x、y為訓(xùn)練數(shù)據(jù)中的輸入、輸出;γ為批量歸一化(BN)的縮放因子;g(γ)為稀疏對比例因子的懲罰函數(shù);λ為平衡因子.

通道剪枝主要是引入BN中的縮放因子γ和稀疏訓(xùn)練中的懲罰項L1.最常見的操作方式是首先在模型中的每個通道都放入一個γ,與通道的輸出結(jié)果進(jìn)行乘積運算,在進(jìn)行稀疏化時將這些比例因子與聯(lián)合訓(xùn)練的網(wǎng)絡(luò)權(quán)重結(jié)合起來進(jìn)行正則化；然后按照縮放因子的大小進(jìn)行排序,設(shè)置一個閾值,根據(jù)閾值將小的縮放因子通道直接裁剪掉,重復(fù)多次進(jìn)行該操作,選取效果較好的一組；最后對網(wǎng)絡(luò)進(jìn)行微調(diào)即可[7].

層剪枝操作是針對模型中每一層的shortcut層的前一個CBL進(jìn)行評價,每層γ的取值按照從大到小規(guī)則進(jìn)行排序,取最小的進(jìn)行層剪枝,每剪去一個shortcut結(jié)構(gòu),相當(dāng)于減掉1個shortcut層和它前面的2個卷積層[8].

在通道剪枝結(jié)束后,將模型再匯總進(jìn)行層剪枝,最后輸出剪枝完的模型.

2.2 知識蒸餾

目前,深度學(xué)習(xí)需要對海量的數(shù)據(jù)進(jìn)行訓(xùn)練來提取特征,這些數(shù)據(jù)數(shù)量較大且冗余,需要消耗大量資源,在沒有實時性要求下訓(xùn)練得到的模型精度較高,但是模型結(jié)構(gòu)笨重.在實際應(yīng)用中,考慮各方面成本,特別是在計算資源方面有所限制,使用一個較大的神經(jīng)網(wǎng)絡(luò)模型有較多局限性.因此,本文采用知識蒸餾模型,讓大模型為小模型進(jìn)行知識遷移.

知識蒸餾的原理是將復(fù)雜的教師網(wǎng)絡(luò)中的dark knowledge遷移到簡單的學(xué)生網(wǎng)絡(luò)中[9],知識蒸餾模型示意圖如圖3所示.

圖3 知識蒸餾模型示意圖

知識蒸餾模型用簡單目標(biāo)來輔助困難目標(biāo)一起訓(xùn)練,soft target來自于教師網(wǎng)絡(luò)的預(yù)測輸出[10],通過在softmax中增加溫度參數(shù)T來實現(xiàn)簡單目標(biāo),公式為:

(2)

式中:qi為函數(shù)的目標(biāo)輸出;T為溫度參數(shù),是從統(tǒng)計力學(xué)的玻爾茲曼分布中借用的概念,當(dāng)溫度T趨向于0時,softmax輸出將收斂為一個one-hot向量,溫度T趨向于無窮時,softmax的輸出則更「軟」;Z為softmax層輸出的分類類別概率.

在訓(xùn)練新模型時,使用較高的T使softmax產(chǎn)生的分布足夠軟,讓新模型(同樣溫度下)的softmax輸出近似原模型;在訓(xùn)練結(jié)束后使用正常的溫度T來預(yù)測.

3 試驗與分析

本文所有的模型訓(xùn)練及相關(guān)試驗均在深度學(xué)習(xí)框架pytorch下實現(xiàn),試驗的硬件配置為:Intel i5處理器,8 G內(nèi)存、NVIDIA GTX 2070super GPU、 PC、Windows 10 操作系統(tǒng).

3.1 數(shù)據(jù)集的制作

國際上公開的海面船舶目標(biāo)檢測數(shù)據(jù)集較少,為評估本文船舶目標(biāo)檢測模型的準(zhǔn)確性,自建船舶目標(biāo)檢測的小樣本數(shù)據(jù)集.通過網(wǎng)上截取船舶視頻片段進(jìn)行視頻的分幀得到一定數(shù)量的船舶圖像,并采用幾何變換中的旋轉(zhuǎn)、裁剪、像素變換中的高斯噪聲等方式進(jìn)行數(shù)據(jù)集的擴(kuò)增,即先對數(shù)據(jù)集采用旋轉(zhuǎn)方式進(jìn)行擴(kuò)增,再隨機(jī)選取一部分圖像進(jìn)行裁切和高斯噪聲操作,將兩次擴(kuò)增的數(shù)據(jù)集集合構(gòu)成總的船舶數(shù)據(jù)集;采用LabelImg標(biāo)定工具對數(shù)據(jù)集進(jìn)行標(biāo)定,得到符合YOLO格式的txt標(biāo)簽文件;將標(biāo)注后的圖像數(shù)據(jù)集中的樣本按照7∶2∶1的比例劃分為訓(xùn)練集、驗證集和測試集,擴(kuò)增后的船舶數(shù)據(jù)集信息如表1所示.

表1 船舶數(shù)據(jù)集信息

3.2 試驗與結(jié)果

首先對船舶目標(biāo)檢測模型進(jìn)行剪枝壓縮,將制作的船舶數(shù)據(jù)集放入YOLO v3模型中進(jìn)行訓(xùn)練；然后對模型進(jìn)行稀疏訓(xùn)練后進(jìn)行通道剪枝、層剪枝；最后對剪枝后的模型進(jìn)行微調(diào),在微調(diào)的過程中對模型進(jìn)行知識蒸餾,使模型得到更好的效果.

表2為不同剪枝系數(shù)的通道剪枝結(jié)果對比.由表2可見,當(dāng)剪枝系數(shù)在0.85時效果最好,經(jīng)過剪枝后模型的參數(shù)量由61 523 734個減少到2 323 692個,減掉約96%,檢測精度由0.931 304降低到0.920 625,下降近1%,總體效果較好.

表2 不同剪枝系數(shù)的通道剪枝結(jié)果對比

通道剪枝后進(jìn)行層剪枝,表3為不同剪枝系數(shù)的層剪枝結(jié)果對比.由表3可見,剪枝系數(shù)增大，檢測精度不斷下降,參數(shù)量下降緩慢,層剪枝沒有通道剪枝效果明顯；剪枝系數(shù)為3時檢測精度最高,與原始檢測精度接近,參數(shù)量沒有太多減少；綜合來看，選取剪枝系數(shù)為3來進(jìn)行層剪枝效果最好.

表3 不同剪枝系數(shù)的層剪枝結(jié)果對比

圖4為模型剪枝前后的檢測效果對比圖.由圖4可見,剪枝后的模型檢測精度與剪枝前相比沒有太多下降,甚至在小目標(biāo)檢測中剪枝后的模型檢測精度比剪枝前有所提升.

(a) 剪枝前

(b) 剪枝后

(c) 剪枝與微調(diào)后

4 結(jié)語

本文采用模型剪枝與知識蒸餾技術(shù)對船舶目標(biāo)檢測模型進(jìn)行壓縮.通過剪枝技術(shù)將模型的參數(shù)量減少了超過90%,大大壓縮了模型體積,使該模型能夠部署到內(nèi)存空間條件不足的硬件設(shè)施上,提升了模型的計算速度,提高了模型的實時性,解決了大部分硬件設(shè)施運行慢的問題.對模型進(jìn)行剪枝后采用知識蒸餾技術(shù)對檢測精度進(jìn)行微調(diào),提高了模型的檢測精度.