語義分割任務(wù)作為計算機(jī)視覺中的基礎(chǔ)任務(wù)之一，其目的是對圖像中的每一個像素進(jìn)行分類。該任務(wù)也被廣泛應(yīng)用于實(shí)踐，例如自動駕駛和醫(yī)學(xué)圖像分割?，F(xiàn)有流行的用于語義分割任務(wù)的模型，例如DeepLab和PSPNet系列，雖然獲得了很好的分割精度，但是所需的算力成本較高。本文考慮使用知識蒸餾（Knowledge Distillation，KD）算法來緩解這個問題。KD作為模型壓縮里的一項(xiàng)重要技術(shù)，其核心思想是將教師模型中學(xué)習(xí)的知識作為監(jiān)督信號來訓(xùn)練學(xué)生模型，使得學(xué)生模型得到更好的性能。雖然現(xiàn)有的語義分割知識蒸餾方法能夠?qū)W(xué)生網(wǎng)絡(luò)的性能進(jìn)行提升，但是這些方法通常是從單張圖像中提取知識，忽略了跨圖像之間的關(guān)系信息也是一種有價值的知識。

最近，地平線-中科院提出了一種新穎的跨圖像關(guān)系型KD方法用于語義分割任務(wù)-CIRKD。該方法嘗試在語義分割任務(wù)中建模pixel-to-pixel和pixel-to-region這兩種對比關(guān)系作為KD的監(jiān)督信號。本文主要從方法介紹以及實(shí)驗(yàn)結(jié)果對提出的CIRKD進(jìn)行講解。

KD作為一種思想已被廣泛應(yīng)用，尤其是圖像分類任務(wù)。但是先前的工作表明直接將圖像分類KD方法直接遷移到語義分割任務(wù)是不可行的，因?yàn)樗鼈儧]有考慮到語義分割的結(jié)構(gòu)化信息。一些經(jīng)典的語義分割KD工作于是去嘗試捕捉到一些上下文信息來建模結(jié)構(gòu)化知識，比如pixel相似度，pixel與region向量的相似度關(guān)系，但是這些方法通常在單張圖像內(nèi)部進(jìn)行信息提取，難以捕捉到更加廣泛的上下文依賴。

于是，我們提出一種跨圖像的知識蒸餾方法來建模圖像間的依賴，從而使得網(wǎng)絡(luò)能夠捕捉到更加豐富的結(jié)構(gòu)化信息。具體地，本工作建模數(shù)據(jù)集中的全部圖像之間建模pixel-to-pixel以及pixel-to-region的對比關(guān)系作為知識。動機(jī)在于一個好的教師模型可以產(chǎn)生更好的全局pixel依賴以及特征空間。CIRKD引導(dǎo)學(xué)生模型去模仿教師網(wǎng)絡(luò)產(chǎn)生的更好的結(jié)構(gòu)化語義關(guān)系，因此提升了語義分割表現(xiàn)。

具體方法

本方法分別從mini-batch和memory的角度對結(jié)構(gòu)化關(guān)系進(jìn)行建模。

1. 基于mini-batch的pixel-to-pixel蒸餾

首先是在batch中建模關(guān)系，將圖像產(chǎn)生的特征圖在spatial維度上進(jìn)行分離，產(chǎn)生pixel-wise的特征向量。針對某一個pixel點(diǎn)的特征向量作為錨樣本，其他pixel點(diǎn)特征向量（包括當(dāng)前圖像和其他圖像）作為對比樣本，可以得到相似度矩陣?；诖朔椒ǎ诮處煻撕蛯W(xué)生端都可以產(chǎn)生對應(yīng)的相似度矩陣，然后通過KL散度進(jìn)行逼近，使得學(xué)生可以學(xué)到教師的相似度矩陣：

一個batch內(nèi)的圖像進(jìn)行兩兩之間的關(guān)系矩陣求取和遷移：

整個過程的示意圖如下所示：

2. 基于memory的pixel-to-pixel蒸餾

相對于基于mini-batch的對比性關(guān)系求取，這里采用一個memory bank來存儲對比向量，使得每一個錨樣本都可以得到充足的對比樣本。針對某一個pixel點(diǎn)的特征向量作為錨樣本，每次訓(xùn)練隨機(jī)從memory bank中采樣大量的pixel特征向量來構(gòu)造對比樣本，通過相乘分別得到教師和學(xué)生的相似度矩陣：

使用KL散度的誤差形式將教師和學(xué)生的pixel-to-pixel相似度矩陣進(jìn)行對齊：

3. 基于memory的pixel-to-region蒸餾

離散化的pixel特征向量不能充分在捕捉到圖像內(nèi)容，因此本方法利用memory對region向量進(jìn)行存儲。region向量產(chǎn)生自對來自相同類別的pixel向量做一個平均。針對某一個pixel點(diǎn)的特征向量作為錨樣本，每次訓(xùn)練隨機(jī)從memory bank中采樣大量的region特征向量來構(gòu)造對比樣本，通過相乘分別得到教師和學(xué)生的相似度矩陣：

使用KL散度的誤差形式將教師和學(xué)生的pixel-to-region相似度矩陣進(jìn)行對齊：

基于memory進(jìn)行KD的整體示意圖如下所示：

最終的訓(xùn)練誤差則是將以上3部分誤差相加在一起：

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

實(shí)驗(yàn)在流行的Cityscapes場景分割數(shù)據(jù)集上進(jìn)行，相比于其他流行的語義分割KD方法，CIRKD獲得了一致最佳的表現(xiàn)。

在一些case study上進(jìn)行可視化，本方法能夠得到最好的分割表現(xiàn)。
審核編輯：彭靜