曠視研究院張祥雨團隊的張培圳研究員深入淺出地為我們介紹了其團隊被 NeurIPS 2021 錄用的論文「用於物體檢測的例項條件知識蒸餾」。本期報告首先簡要回顧了用於影象識別、物體檢測的知識蒸餾技術的發展歷程,介紹了論文的研究動機、模型架構、訓練技巧以及實驗結果。
本文整理自青源 LIVE 第 31 期,影片回放連結:https://event.baai.ac.cn/activities/204
主講丨張培圳
整理丨熊宇軒
審校丨李夢佳
1 背景知識
早在 NIPS 2014 上,Hinton 等人在論文「Distilling the knowledge in a Neural Network」中就開始研究用於影象分類的知識蒸餾技術。根據知識源的不同,我們可以將用於影象分類的知識蒸餾技術分為「基於 Logits 響應」、「基於中間特徵圖啟用」和「基於關係建模」的知識蒸餾。
如上圖所示,基於響應的知識蒸餾往往會在弱監督場景下計算預訓練教師網路和學生網路用於分類的 Logits 的 KL 散度;基於啟用的知識蒸餾會計算教師網路和學生網路中間特徵圖的 MSE;基於關係建模的知識蒸餾則會針對教師網路和學生網路對每一個 mini-batch 中圖片經過全域性平均池化後的特徵矩陣計算 MSE 等距離。
面向物體檢測任務的知識蒸餾技術大多屬於基於中間特徵圖啟用的知識蒸餾。近年來,一些面向工程化的工作也會嘗試系統地使用上述知識蒸餾方式的組合(例如,FBKD、GID、DeFeat、G-DetKD,等)
相較之下,分類模型需要處理的影象往往包含單個物體,待分類的物體位於影象的中央;而檢測模型需要處理的影象中往往包含多個物體,這些物體從屬於各種類別、分佈於影象的各個位置,影象中還包含大量冗餘的背景資訊。
在針對中間特徵圖進行蒸餾時,蒸餾特徵的位置是十分關鍵的,我們通常會利用以下三類拓撲資訊:
(1)前景先驗:給定訓練圖片中物體的邊界框 BBOX,我們將其對映到中間特徵圖上,對該區域進行知識蒸餾;對於基於 anchor 的物體檢測任務,重點對與 groundtruth 重疊比例較高的 anchor 區域進行蒸餾;對於多階段物體檢測任務,針對 proposal 進行蒸餾;
(2)前景和背景先驗:將前景、背景分開蒸餾/在全圖上使用注意力機制;
(3)基於度量的區域:根據設計出的度量方式找出有效的蒸餾資訊(例如,教師網路與學生網路的差異)。
如上圖所示,藍色和棕色的實線方框代表真實的物體邊界框;黑色的方框代表物體附近的 anchor;藍色、棕色、紅色的虛線方框代表多階段檢測模型提出的 proposals,可以透過非極大抑制等手段對其進行修正。
在考慮背景先驗時,我們可以透過二值掩膜將前景和背景分隔開,進而分別進行蒸餾。此外,我們可以在全圖中分別在空間和通道上使用注意力機制,我們可以分別對教師和學生網路計算空間維度和通道維度上的注意力圖,進而對這些注意力兩兩進行知識蒸餾。
在利用教師網路、學生網路評分差異進行知識蒸餾時,假設教師網路和學生網路預測頭的預測特徵圖的維度相同,影象上的某個畫素點鄰近於某個 anchor 的預測框,模型會計算出該預測框從屬於某一類的置信度以及物體迴歸的位置編碼資訊。在典型的情況下,教師網路可以很好地預測出類別和邊框位置資訊,而學生網路對於某些類別物體的預測效能較差。因此,我們可以基於度量學習的方式,重點對學生網路學習效能較差的類別進行蒸餾。
以 ADL 模型為例,其關鍵之處在於計算自適應蒸餾權重(ADW)。如上圖所示,紅色圓圈代表教師網路和學生網路在某個空間位置上預測評分的 KL 散度。此外,在 GID 模型中,作者定義了一種通用例項(General Instance)代表教師網路或學生網路的 proposal,這兩種網路對proposal 會產生預測得分的差異,GI Score 為針對所有類的得分差異的最大值,該值越大說明不確定性越強。
基於前景/背景的建模方法非常直接,符合人類的認知習慣,但是這些方法獲得的知識可能是次優的。基於度量的方法相對來說較為靈活,但是往往並不直觀,並且存在多種可能的度量方法。
2 用於物體檢測的例項條件知識蒸餾
為了解決上述問題,曠視研究院張祥雨組提出了「用於物體檢測的例項條件知識蒸餾」,在保證模型具有直觀性、可解釋性的同時,兼具了基於度量的蒸餾方法的靈活性。模型的整體架構如上圖所示。
具體而言,該方法利用人類標註的例項作為條件,來查詢教師網路的特徵,從而檢索出有用的知識,用於蒸餾學生網路的特徵圖。
如上圖所示,左側為 Query 的特徵構造過程,右側為抽象的建模過程。從抽象建模過程開始,我們可以將圖中物體的 one-hot 類別資訊和例項的位置嵌入透過多層感知機編碼到 Query 中。這樣一來,每個例項的特徵都可以被表示為一個低維向量。
我們利用上述例項標註在教師網路的特徵圖上提取有效資訊。每個特徵圖是一個三維(H*W*D)的張量,我們可以將其在空間維度上展平,得到特徵
接著,我們透過「鍵-值」形式的注意力,將物體標註和保留了空間-拓撲資訊的特徵影象素之間的關係例項化,從而得到了「例項感知注意力掩模」。這樣一來,我們就可以衡量物體和影象位置的關係。其中,基於 Query、Key、Value 的自注意力的計算正規化與原始的 Transformer 架構中的計算正規化相同。
為教師網路的特徵圖,i 為圖中第 i 個物體的索引,j 為多頭注意力機制中的第 j 個注意力頭。
為了保證掩模與待檢測物體的位置相關,我們設計了面向檢測任務的輔助任務引導對掩模的學習。具體而言,我們利用注意力機制的輸出作為識別、定位輔助任務的監督訊號。值得注意的是,輔助任務不會更新學生網路的引數。在識別任務中,我們首先根據資料集先驗取樣出一些偽樣本。對於識別出的真實樣本,我們進一步對其位置進行定位。
如上圖所示,在公式(10)中,輔助損失中的第一項是針對識別任務的二值交叉熵損失,第二項為對真實樣本的定位迴歸損失。其中,
為由多層感知機和 Sigmoid 函式構成的分類頭,當查詢的標註物體為真時
然而,注意力機制的輸出依賴於注意力掩模和教師網路特徵圖,而注意力掩模又依賴於標註。因此,透過公式(10)懲罰注意力機制的輸出和之間的分類誤差,往往會使網路學習到一個平凡解,從而使「例項感知注意力掩模」失效。
為此,我們引入了資訊丟失(information dropping)技術,破壞例項標註的資訊。對於迴歸任務,我們對物體邊界框的中心進行抖動,並且對尺度資訊進行量化。這樣一來,我們就會強迫模型從教師網路特徵圖上學習到更多的資訊。
因為資訊丟失和定位迴歸損失之間有聯絡,資訊丟失能強迫從教師網路特徵圖上學習到更多的資訊。我們的定位任務不是像一般檢測器那樣去迴歸物體框,而正是去迴歸抖動的物體框中心和抖動前的框中心的這個抖動距離。因為 本身是基於抖動後的框資訊編碼的,它如果想把上面這個任務做好,降低其損失,必須從教師特徵圖上面去感知到這個真實物體中心。
整體而言,本文所提出的框架的損失函式由蒸餾損失、學生網路的檢測損失、輔助任務損失構成,其中
3 實驗結果
與現有的知識蒸餾方法相比,本文提出的方法在 COCO 資料集上取得了較優的效能。在使用 RetinaNet 作為物體檢測網路,採用多尺度訓練時,本文提出的方法可以將學生網路(ResNet-50)的物體檢測 AP 從 37.4 提升到 40.7,這一數字甚至超過了教師網路(ResNet-101)的 40.4 的 AP。在訓練過程中進行的,涉及物體Groundtruth資訊的例項建模為學生從單純的教師網路特徵圖資訊以外獲得更好訓練,提供了可能。
在其它的檢測和分割模型上,本文提出的知識蒸餾方法也取得了優秀的效能表現。如上圖所示,繼承策略(Inheriting Strategy)是一種非常有用的訓練技巧:當教師網路和學生網路的Neck部分和檢測頭部分的架構相同時,學生網路就可以繼承上面兩個部分的引數進行良好的初始化,從而進一步提升效能或快速收斂。
在使用 EfficientNet-B0、MobileNet-V2 等面向移動裝置的主幹網路時,本文提出的知識蒸餾方法可以大大提升物體檢測的 AP。
Transformer 解碼器的注意力頭在原圖中的啟用視覺化結果如上圖所示。實驗結果表明,注意力頭可以有效注意到例項周圍的資訊。
4 討論
在整體的計算過程中,計算量主要體現在教師網路的預訓練和學生網路本身的訓練開銷,而本文加入的輔助任務和知識蒸餾的計算開銷相對較小。在知識遷移方面,本文提出的方法可以直觀地檢索影象中有意義的內容,引入了類似基於度量的蒸餾方法的靈活特性,具有一定的可解釋性。本文采用的繼承策略訓練技巧可以有效地加快收斂。
5 結語
在本文中,曠視研究院提出了一種面向物體檢測的新型知識蒸餾框架,該框架可以學習特徵解碼器模組從而檢索有用的知識,設計了輔助任務來最佳化解碼器。作者設計了基於例項的條件,從而適應目標檢測任務,其中 Query 為標註的例項資訊,Key 和 Value 為教師網路的表徵知識。
GAIR 2021大會首日:18位Fellow的40年AI歲月,一場技術前沿的傳承與激辯
2021-12-10
致敬傳奇:中國並行處理四十年,他們從無人區探索走到計算的黃金時代 | GAIR 2021
2021-12-09
時間的力量——1991 人工智慧大辯論 30 週年紀念:主義不再,共融互生|GAIR 2021
2021-12-12
論智三易,串聯通訊,貫通邊緣,演進認知,匯於機器:聽五位IEEE Fellow暢談AI未來 | GAIR 2021
2021-12-25