基于掩膜的目標檢測模型蒸餾方法

陳宇軒，張本亭，濮約剛，張明慶

(中國航天科工集團第二研究院 七〇六所，北京 100854)

0 引 言

在本世紀初，隨著深度學習回歸大眾視野并成為研究人員關注的重點，目標檢測算法也得到進一步發展。近些年來，隨著深度學習的進一步發展，學術界中的神經網絡層數越來越深，參數量越來越大，模型性能也在不斷提升，需要的算力和內存資源也越來越大，對計算設備的需求也與日俱增。需求是巨大的生產力，目前的GPU和AI芯片也如雨后春筍一般飛速發展。目標檢測網絡也享受到了計算設備發展的紅利，在模型訓練逐漸使用更深層的特征提取網絡來提取特征圖片特征，檢測精度一次又一次取得突破。在工業界，大量目標檢測算法被應用于自動駕駛、人臉識別等領域，將目標檢測算法部署在移動端往往受到算力和內存空間的限制。在實際部署時，使用精度高的大模型部署在移動端平臺并不能達到實時檢測的要求，而使用檢測速度快的小模型往往檢測精度不夠。為了解決這一對巨大的矛盾，使得檢測算法可以平穩落地部署，誕生了許多行之有效的方法，知識蒸餾便是其中之一。Hinton提出像教師和學生一樣使用大網絡輸出的概率分布作為教師網絡來指導小網絡進行訓練，使得學生網絡的輸出分布最終接近于教師網絡，這種方法被命名為知識蒸餾。隨后出現了大量關于知識蒸餾策略的研究，盡管知識蒸餾已經在分類等領域取得巨大的成就，但對于目標檢測任務的蒸餾始終是一個難點。其原因在于目標檢測相比于分類任務更為復雜，需要傳遞更多的知識，此外目標檢測中的背景類會對檢測結果產生一定影響。目前針對目標檢測的知識蒸餾方法，為了避免背景類對蒸餾造成影響，通常會采用過濾背景進行蒸餾的方法[1]，或者將背景和前景區域分開進行蒸餾[2]；然而，最近的研究表明，物體的背景類中也包含重要的信息，僅對前景進行蒸餾可能會忽視掉這些信息，因此在蒸餾時不應簡單的將背景與前景分開[3]。為此，本文提出一種基于掩膜的目標檢測模型蒸餾方法，該方法同時對背景類和前景類展開蒸餾，通過掩膜給不同的區域賦予相應的權重，對于權重更大的位置在蒸餾時給予更大的關注。

1 基于目標檢測的知識蒸餾方法

在目標檢測上的知識蒸餾需要克服更多的困難。首先，目標檢測任務相比于其它任務本就更加困難，除了需要獲得目標的類別，還需要了解目標的定位，這就意味著教師網絡需要傳遞更多的知識給學生網絡；此外，在目標檢測中，除了物體類還包括背景類，教師網絡中所傳遞的背景知識和前景知識如果不加以區分，學生網絡會從中得出誤導性的結論。

Guobin Chen等在目標檢測中應用知識蒸餾壓縮模型，使用二階段目標檢測器Faster R-CNN對學生網絡進行蒸餾，提出了用新的Loss來處理上述的問題和挑戰[4]。在對主干網絡的蒸餾上，注意到使用教師網絡的中間表示作為提示(Hint)可以幫助訓練過程并提高學生網絡的最終表現。

Wang等在隨后的研究中發現，提示部分的損失僅僅只對全局信息進行蒸餾，而全局信息中必定帶有大量的背景信息，針對這些魚龍混雜的信息進行相同的蒸餾無法達到最好的效果。為此，Wang等提出了Fitnet對此進行改進，通過在蒸餾過程中加入掩膜除去教師網絡中的背景，學生網絡進行蒸餾時只關注于目標周圍的特征情況[1]。圖1展示了掩膜區域的生成過程。

對于每一個目標的真實標簽都會計算其與教師網絡生成的所有檢測框的IoU值，對于其中大于閾值F的所有檢測框，將其進行組合作為該目標的掩膜區域。通過掩膜區域屏蔽背景信息的干擾，該部分的損失為在掩膜區域監督下學生網絡向教師網絡學習產生的損失，具體計算方式如式(1)所示

(1)

式中：NP為檢測目標中正例的數量，I為Fine-Grained Imi-tation Mask的掩碼Mask，fadap(·)為特征適應層函數。整個蒸餾的最終損失如式(2)所示

L=Lgt+λLimitation

(2)

Limitation為學生網絡根據掩膜向教師網絡學習的損失。Limitation為學生網絡向真實標簽學習產生的損失，具體計算方式沿用了Hinton的計算方法。

Linfeng Zhang等對Wang等的方法作出了改進，設計出了基于注意力引導機制的掩膜用于目標檢測的蒸餾，通過加入注意力機制分別在通道和空間維度上產生通道和空間注意力，并鼓勵學生網絡模仿教師網絡的注意力，與Wang等的方法相比，這種方法生成的掩膜增加了通道掩膜，并且可以做到像素級別的掩膜[5]。其從目標特征到注意力的映射方式如式(3)、式(4)所示

(3)

(4)

其中，A表示目標檢測模型中主干的特征，C，H，W分別表示其通道數，高度和寬度，GS為通道注意力，GC為空間注意力。

空間注意力掩膜和通道注意力掩膜的生成方式如式(5)、式(6)所示

MS=HW·softmax((GS(AS)+GS(AT))/T)

(5)

MC=C·softmax((GC(AS)+GC(AT))/T)

(6)

其中，H，W，C分別表示目標的高度，寬度和通道數，AS和AT分別為學生網絡和教師網絡的由主干網絡提取出的特征。

最終的注意力蒸餾損失由兩部分構成，注意力轉移損失和注意力掩膜損失，計算方式如式(7)所示

LAGD=LAT+LAM

(7)

式中：LAT為注意力轉移損失，用于鼓勵學生網絡模仿教師網絡的空間注意力和通道注意力，計算方式如式(8)所示

LAT=L2(GS(AS)，GS(AT))+L2(GC(AS)，GC(AT))

(8)

LAM為注意力掩膜損失，用于鼓勵學生網絡通過空間注意力掩膜和通道注意力掩膜學習教師網絡的特征，計算方式如式(9)所示

(9)

Wang和Zhang的方法使用掩膜過濾了圖像中的背景信息過濾，學生網絡只需要向教師網絡的前景部分學習，但是該類方法忽略了對背景中知識的蒸餾，Jianyuan Guo等指出，物體和背景之間有很強的關系，僅通過背景區域特征提取，學生網絡也可以顯著提高學習能力，甚至可以接近于前景區域提取的蒸餾結果[3]。

目前對于目標檢測的蒸餾研究，通常選擇直接屏蔽背景以去除背景對于檢測的干擾，但是忽略了背景中包含的重要信息，此外，在對前景進行學習時并沒有區分其重要程度。同時，之前的方法大多數是針對二階段的目標檢測方法，而實際工業應用中大多采用一階段目標檢測方法，因此此前所提出的方法并不能直接應用于實際場景。為此，本文針對目前蒸餾的主要問題作出改進，基于以往的掩膜設計方法提出用于平衡教師網絡前景信息與背景信息的帶權值的掩膜。

2 基于掩膜的知識蒸餾方法設計

Fitnet中指出，教師網絡可以提取出輸出豐富的全局信息和目標的定位信息，可以良好地體現教師網絡的泛化能力[1]，通過比較中間層學生網絡與教師網絡特征的方法，相當于引入一個助教，引導學生網絡學習教師網絡中間特征層的輸出，通過學習這部分的知識，可以獲得接近于教師網絡的泛化能力。然而，在目標檢測中，教師網絡中有許多冗余元素，學生網絡并不知道哪一部分的特征更值得學習，如果只是平等地學習每一部分知識，腦容量不夠的學生網絡并不會取得太好的學習效果。并且，教師網絡產生的特征往往比較復雜，對于這一類教師網絡產生的知識，需要讓學生網絡知道學習的重點，才能更好地學習教師網絡的知識，做到有的放矢。此外，對于目標檢測，前景和背景之間傳遞的信息量并不一致，在圖像中占據大部分的背景包含的知識相比于前景顯得不值一提，如果一味將背景知識和前景知識平等地教授給學生網絡，會影響最終學生網絡學習的結果。因此許多方法選擇了過濾背景后進行蒸餾，但如果僅是針對前景進行蒸餾，可能會忽略背景中潛在的信息區域，并且來自背景的一些信息也可以幫助學生網絡學習教師網絡的泛化。為此，本文提出的基于帶權值的掩膜的蒸餾在教師網絡傳遞知識的過程中，會對前景與背景的重要程度進行區分，在學習前景知識的同時，保留了部分背景知識幫助學生網絡學習教師網絡的泛化；同時對前景知識的重要程度進行區分，讓學生網絡知道學習的重點，從而更好地學習教師網絡的知識。

2.1 蒸餾架構

目標檢測網絡主要由主干網絡Backbone、脖頸網絡Neck、檢測頭Head等幾部分組成，其中主干網絡負責圖片特征提取，提取不同尺度、不同感受野下、不同類別的目標特征，用以目標檢測；脖頸網絡部分主要是將主干網絡上的多個層級的特征圖進行融合拼接，獲得來自不同感受野的信息，增強網絡的表達能力。檢測頭Head主要用于預測目標的種類和位置。本文提出的蒸餾方案使用帶有權重的掩膜對特征提取層進行蒸餾，使學生網絡對于教師網絡傳遞的知識進行有側重點的學習；同時在脖頸網絡上采用Object Relation Module通過教師網絡和學生網絡的全局特征關系展開蒸餾，在一次次訓練中，使學生網絡更加關注教師網絡的全局上下文中的知識傳遞，彌補了掩膜方案中不能傳遞全局知識的缺陷。相比于目前常用的蒸餾架構，可以有效提升蒸餾效果。整體蒸餾架構如圖2所示。

圖2 基于帶權值的掩膜蒸餾架構

圖片分別輸入教師網絡和學生網絡，首先在骨干網絡層對教師網絡與學生網絡提取的特征展開蒸餾，根據教師網絡與真實標簽計算出學生網絡的帶權掩膜，使用得出的掩膜指導學生網絡向教師網絡學習，權重高的位置的損失值更大，將得到更多的關注，并計算該部分Weighted Background &Foreground 損失；接著對全局關系展開蒸餾，并分別計算出教師網絡和學生網絡經過Object Relation Module得出的新特征，以此來計算出教師網絡與學生網絡之間的全局特征關系間的差距，并使學生網絡學習教師網絡中的全局關系；最后通過凍結教師網絡權重，使用反向傳播算法更新學生網絡的權重，在多次訓練中，使學生網絡逐漸掌握教師網絡的特征知識。

2.2 掩膜計算

帶權重的掩膜計算過程如圖3所示，首先，將圖像輸入教師網絡生成多個目標檢測預測框，同時得到教師網絡的置信圖Confidence-map。之后，由教師網絡輸出的物體的預測框與真實標簽，篩選出包含目標檢測物體的前景掩膜圖，同時得到不包含目標的背景掩膜圖。

圖3 帶權值的掩膜生成步驟

為了計算出目標檢測背景的重要程度，在檢測中使用置信度對其加以衡量，置信度即是用來表示輸出的框中是否有物體的自信程度，置信度的值越大，即表明該處有物體的可能性越大。置信度用來判斷邊界框內的物體是正樣本還是負樣本，大于置信度閾值的判定為正樣本，小于置信度閾值的判定為負樣本即背景。具體而言，距離物體越遠處的背景置信度越小，而更靠近前景的部分置信度則更大。本文根據教師網絡輸出的置信度，用以計算當前目標檢測中背景的重要程度，具體做法見式(10)

(10)

式中：N為樣本中負例的個數。

計算出背景的重要程度之后，為了讓學生網絡知道前景部分中學習的重點，需要區分教師網絡前景不同區域的重要程度，一個直觀的想法是越靠近檢測中心部分的知識越重要，因此，本文根據之前計算出的Background-threshold，以及由預測框與真實標簽生成的前景掩膜圖，計算出帶權重的掩膜，具體做法見式(11)

(11)

式中：BT為Background-threshold，XC與YC為檢測中心點位置，S為前景圖中距檢測中心點的最遠距離。根據式(11)獲得帶權重的掩膜圖，在計算損失時，權重更大的位置會受到更多關注，在保留背景區域的前提下使學生網絡更加關注前景區域中重要位置的學習，從而更好地擬合教師網絡產生的特征。

2.3 基于全局特征關系的蒸餾

在卷積神經網絡中，卷積層在局部鄰域內構建像素關系，可以提取出局部特征，但是在遠距離位置之間傳遞信息存在困難。現有的目標檢測算法都把圖像中的物體看成一個個獨立的個體去檢測，并沒有利用不同物體間的相關性。Cao等指出，提取對視覺場景的全局理解，有利于廣泛的識別任務[6]。針對以上問題，本文采用了Object Relation Module來利用圖像中物體之間的相互關系或者圖像內容[7]，在獲取了教師網絡和學生網絡的特征關系后，通過設計對特征關系的蒸餾使學生網絡學習教師網絡中這一類全局關系，以補充對于全局關系的蒸餾。

在Object Relation Module中，每一個目標的最終特征由本身提取出的特征與他收集到的與其它特征的關系fR組成，其計算方式如式(12)所示

(12)

(13)

(14)

式中：分子為第m個物體和第n個物體的外觀特征投影到低維空間的點積，用投影的重合程度來表示其影響程度，分母為降維后的維度

(15)

式中：ε表示將特征映射到高維空間，之后通過WG進行變換。

通過以上方式收集到物體n的一類關系特征fR(n)， 通常會收集16組并將其Concat到一起，最后與原來的外觀特征相加得到最終的特征。計算方式如式(16)所示

(16)

在得到教師網絡的最終特征fA(T) 與學生網絡最終特征fA(S) 后，本文通過計算fA(T) 與fA(S) 之間L2距離的平方來衡量教師網絡與學生網絡之間的差距，計算方式如式(17)所示

(17)

2.4 損失設計

整體蒸餾損失主要分成3個部分，具體見式(18)

Ltotal=(1-α-β)Lhard(S，LABEL)+αLweighted B&F(S，T)+βLObject Relation(S，T)

(18)

如式(18)所示，Ltotal為蒸餾總損失，Lweighted B&F為學生網絡與教師網絡間的帶權特征蒸餾損失，計算方式見式(19)，LObject Relation為學生網絡與教師網絡之間全局特征關系蒸餾損失，計算方式見式(17)，α和β為超參數，用以平衡三者之間的權重，為鼓勵學生網絡更多向教師網絡學習而不是通過True Label進行自學，在本文中均設置為0.45。Lhard為目標檢測訓練時向Hard Target即真實標簽學習的損失，計算方式沿用了Hinton方法中Hard損失的設計。

Hinton在計算蒸餾損失時將損失分為了衡量由學生網絡輸出與真實標簽之間差距的Hard Loss，以及衡量由學生網絡輸出與教師網絡輸出差距的Soft Loss部分，在本部分中，衡量學生網絡輸出與教師網絡輸出差距的公式如式(19)所示

(19)

式中：N為檢測中總像素點的個數，s為學生網絡的引導特征圖，t為教師網絡的引導特征圖，fadap(·)為自適應層函數，W，H，C表示輸出特征的寬，高和通道數，Weightedmaskij為前文求出的帶權重的掩膜圖。

3 知識蒸餾實驗結果

3.1 實驗環境

采用本文所提出的蒸餾方法的實驗部分在同一臺服務器上完成。實驗環境見表1。

表1 蒸餾訓練實驗環境

服務器的硬件條件如下：CPU為8核Intel Xeon處理器、內存為32 G、使用的GPU為英偉達的Tesla T4，內存16 G。

軟件環境如下：實驗過程中實驗ubuntu18.04操作系統進行實驗，使用python3.9版本作為編程語言，使用的計算加速庫為CUDA11.6與cuDNN8.6，蒸餾訓練框架為pytorch，版本為1.12，采用了opencv-python庫對圖像進行處理。

3.2 數據集與評價指標

本文的知識蒸餾實驗主要基于COCO數據集，COCO數據集中提供的圖像包括有91類目標，其中有80類物體和1類背景類，擁有超過30萬張圖像，其中20萬張圖像有標注，共擁有超越200萬張圖片標注，整個數據集中的物體超過150萬個。

mAP(mean average precision)，意為平均類別準確率，通過計算每一個類別的準確率后，對其相加取平均值得出。

3.3 實驗結果

算法的蒸餾效果如圖4所示，與蒸餾前的Yolov3-Tiny模型輸出結果作為對比，可以看出，蒸餾后的模型效果明顯好于蒸餾訓練前的Yolov3-Tiny模型，針對之前學生網絡中漏掉的檢測框，在進行蒸餾訓練后，提高了檢測能力的學生網絡可以檢測出來，包括一些有遮擋物的回歸框或者尺寸較小的物體，例如圖4(b)中，蒸餾后的網絡檢測出了小物體交通信號燈Traffic Light，以及圖片邊緣的人物，而圖4(f)并沒有檢測出，同時檢測產生的錯誤也更少，圖4(h)中，Yolov3-Tiny將Car檢測成了Bus，并且圖片中存在大量的小目標無法分出類別，圖4(d)蒸餾后的網絡可以檢測出之前無法檢測出的小物體的類別，雖然在部分目標的判斷中出現錯誤，但是檢測框的準確程度與檢測精度相較之前有所提升。在蒸餾后尺寸沒有太大變化的前提下，提升了原有模型的檢測能力。

圖4 蒸餾效果展示

為了更直觀地感受本文蒸餾算法帶來的提升。本文繪制了基于Yolov3-Tiny目標檢測算法在蒸餾訓練時的Object Relation Loss和Weighted Mask Loss的變化曲線，如圖5所示，可以看出，訓練初期教師網絡與學生網絡有著較大差距，隨著訓練代數的增加，教師網絡與學生網絡之間差距逐漸縮小，損失最終趨于平穩。基于Yolov3-Tiny目標檢測算法在蒸餾訓練時的平均類別準確率、召回率的變化曲線如圖6所示，可以看出，蒸餾后的模型mAP相比于未蒸餾模型明顯提升，蒸餾之后模型的性能明顯要優于未蒸餾模型。

圖5 蒸餾訓練損失變化曲線

圖6 訓練精度及召回率變化曲線

為了對比本文提出的蒸餾方案的效果，采用了Fine-Grained知識蒸餾方案作為對比，采取mAP作為檢測指標，mAP50、mAP75表示在采用IoU=0.5標準以及IoU=0.75標準下的mAP值，mAPs、mAPm、mAPl用來統計模型對于小目標、中等目標、大目標的檢測能力。如表2所示，可以看出，Yolov3模型在COCO數據集上的mAP@IoU=0.5達到了56.1%，可以用來當作教師網絡訓練Yolov3-Tiny模型，以Fine-Grained方法訓練后得到的模型在COCO2017數據集上的mAP@IoU=0.5達到了40.2%；以本文Weighted Mask方法訓練后得到的模型在COCO數據集上的mAP@IoU=0.5達到了42.5%；以本文Weighted Mask方法加上Object Relation Module對全局關系進行蒸餾的mAP@IoU=0.5達到了43.6%，相比Yolov3-Tiny模型提升了76.5%，并且無論是大目標的檢測(mAPl)、中等目標的檢測(mAPm)還是小目標的檢測(mAPs)，該方法都擁有更好的效果。此外，本文統計了Yolov3-tiny、Yolov3及蒸餾訓練后模型的FPS及模型參數量，詳情見表3。

表2 YOLO模型不同蒸餾方法下精度

表3 YOLO模型不同蒸餾方法FPS及模型參數量

經過蒸餾后的學生網絡精確度與教師網絡相當，推理效率FPS得到顯著提升，但是參數量指標卻大幅下降，蒸餾后得到的模型擁有接近于教師網絡的檢測精度與學生網絡的推理速度，達到了實際部署的要求。

4 結束語

本文針對目前目標檢測算法背景類和前景類出現的頻率和重要程度不一致，難以蒸餾的特點提出用于平衡教師網絡前景信息與背景信息的帶權掩膜，通過對前景與背景的重要程度進行區分，在學習前景知識的同時，保留了部分背景知識幫助學生網絡學習教師網絡的泛化；同時對前景知識的重要程度進行區分，讓學生網絡知道學習的重點，從而更好地學習教師網絡的知識。

在COCO2017數據集上進行蒸餾訓練，以Yolov3作為教師網絡對Yolov3-Tiny模型進行蒸餾，結果顯示，在模型參數量沒有提升的前提下檢測精度從24.7%提升到43.6%，提升了76.5%。