FRKDNet：基于知識蒸餾的特征提煉語義分割網(wǎng)絡

蔣詩怡，徐 楊，2*，李丹楊，范潤澤

（1.貴州大學 大數(shù)據(jù)與信息工程學院，貴州 貴陽 550025；2.貴陽鋁鎂設計研究院有限公司，貴州 貴陽 550009）

1 引 言

語義分割的目的是為輸入圖像的每個像素進行分類，也是計算機視覺中一項基礎且富有挑戰(zhàn)性的任務，目前已廣泛應用到多個領域，如自動駕駛［1-2］、行人分割［3］、遙感監(jiān)測［4-5］等。近年來，由于深度學習在計算機視覺領域的成功，基于卷積神經網(wǎng)絡（Convolutional Neural Network，CNN）的方法大幅提高了語義分割的準確性。然而，基于CNN 的語義分割算法通常具有較高的計算成本，這也限制了網(wǎng)絡模型在實際中的應用與部署，特別是在對效率要求較高的現(xiàn)實任務中。

為了解決上述問題，研究人員探索了網(wǎng)絡輕量化的方式。Paszke 等人［6］重新設計了網(wǎng)絡架構降低模型參數(shù)量，提高了推理速度。Zhao 等人［7］提出級聯(lián)特征融合單元和級聯(lián)標簽指導策略整合中等和高分辨率特征，提煉了粗糙的語義預測圖。盡管研究人員已經設計了優(yōu)秀的網(wǎng)絡來減少計算成本，但很難在精度和模型大小之間達到令人滿意的效果。針對這一問題，我們采用知識蒸餾（Knowledge Distillation，KD）［8］的策略，通過在訓練過程中蒸餾較大模型（教師網(wǎng)絡）的知識來提高高效輕量模型（學生網(wǎng)絡）的性能。

KD 作為一種模型壓縮的方法，最初用于圖像分類任務中，其能夠顯著簡化繁瑣的模型。由于KD 的優(yōu)勢，一些語義分割方法使用KD 來減小模型的大小。這些方法讓學生網(wǎng)絡從教師網(wǎng)絡學習像素特征和類內特征的變化，具有代表性的是類內特征方差蒸餾（Intra-class Feature Variation Distillation，IFVD）［9］，側重于將類內特征的差異從教師網(wǎng)絡轉移到學生網(wǎng)絡。通道知識蒸餾（Channel-wise Knowledge Distillation，CD）［10］強調在每個通道中提取最重要的區(qū)域。值得注意的是，語義分割是一項包含各種類別的像素級類別預測任務，因此在語義分割中，特征空間中的類間距離是普遍存在的。教師網(wǎng)絡借助眾多參數(shù)和復雜的網(wǎng)絡結構，在特征空間中具有較強的分類能力和較大的類間距離。但以往方法仍存在一定的問題：（1）以往的語義分割KD 方案忽略了將教師網(wǎng)絡特征空間中的類間距離與類內距離轉移到學生網(wǎng)絡；（2）由于學生網(wǎng)絡結構簡單，參數(shù)量較少，學生網(wǎng)絡不能像教師網(wǎng)絡那樣編碼豐富的位置特征信息；（3）教師網(wǎng)絡結構復雜，參數(shù)量大，輸出內容包括前景知識內容和背景噪聲，對學生網(wǎng)絡來說偽知識的判別也是知識蒸餾過程中的一大難題。為了解決上述問題，我們設計了基于知識蒸餾的特征提煉語義分割模型（Feature Refine Semantic Segmentation Network based on Knowledge Distillation，F(xiàn)RKDNet），首先設計了一種特征提煉方法（Feature Refine，F(xiàn)R），來分離蒸餾知識的前景內容與背景噪聲，在特征提煉的幫助下過濾偽知識，將更準確的特征內容傳遞給學生網(wǎng)絡。其次考慮在特征空間的隱式編碼中提取類間距離與類內距離，從而得到相應的特征坐標掩碼，分離相應的特征位置信息，學生網(wǎng)絡通過模擬特征位置信息來最小化與教師網(wǎng)絡特征位置的輸出從而提高學生網(wǎng)絡的分割精度。

2 本文模型

2.1 網(wǎng)絡結構

本文提出一種基于教師-學生模型的知識蒸餾模型，總體模型架構如圖1 所示。整體架構由教師網(wǎng)絡、學生網(wǎng)絡、特征提煉模塊、類內距離計算模塊與類間距離計算模塊組成。對于教師網(wǎng)絡和學生網(wǎng)絡，我們使用經典的編碼器-解碼器結構提取圖像特征并輸出相應的像素分類概率分布特征圖。知識蒸餾一般需要一個強大的教師網(wǎng)絡，而要獲得更詳細的圖像特征，就需要一個足夠深的教師網(wǎng)絡。具體而言，我們將圖像分批輸入到教師網(wǎng)絡和學生網(wǎng)絡中進行訓練，獲得各自的輸出特征圖。針對以往知識蒸餾方法中教師網(wǎng)絡傳遞的知識復雜、特征與噪聲混合的情況，為了排除偽知識對蒸餾的影響，減少背景噪音，我們設計了知識提煉模塊將教師網(wǎng)絡輸出的特征圖中的顯著知識與背景噪聲分離開。在分離過程中，通過計算每個像素分類特征與當前所屬類別的信息熵差值并保留信息量大的顯著信息，從而得到新的特征提煉矩陣，并分別與教師網(wǎng)絡和學生網(wǎng)絡的原始輸出特征相乘得到含有顯著信息的特征圖。針對以往方法中忽略特征內部特性，不同類物體之間差別不明顯，同類物體之間特征不突出的問題，通過提出的類間距離計算模塊與類內距離計算模塊分別計算教師網(wǎng)絡與學生網(wǎng)絡的蒸餾損失，同時計算學生網(wǎng)絡的原始輸出與實際標簽的損失函數(shù)，合并作為整個模型的損失函數(shù)一起反向傳播。

圖1 網(wǎng)絡模型結構Fig.1 Network structure of model

2.2 特征提煉模塊

特征提煉是一種過濾無用信息，提煉顯著信息的過程。在教師網(wǎng)絡蒸餾的知識中通常含有一些背景噪音，通過特征提煉選擇顯著且豐富的知識進行蒸餾與距離計算，從而減少在蒸餾與知識傳遞的過程中噪聲的影響。文獻［8］表明，可以引入與教師網(wǎng)絡相關的軟目標來引導學生網(wǎng)絡訓練。對于語義分割任務來說，每一個像素都會輸出一個像素分類概率，基于知識蒸餾的語義分割任務也是最小化像素間的概率誤差，因此我們認為網(wǎng)絡特征圖中像素對應的值與所屬類別對應的比例即為當前像素的類別貢獻，貢獻越大的即為顯著特征。將原始網(wǎng)絡輸出的特征圖中的像素值轉換為對應的類別預測概率，從而計算該像素點對整個類別預測的貢獻值，其中預測貢獻值為該像素值除以整個類別像素預測總和，如式（1）所示：

其中p(x)為該像素的輸出值。對于教師網(wǎng)絡的輸出特征Ft∈RC×H×W，其中C為通道數(shù)，H與W為空間維度，故先通過Softmax 函數(shù)將特征圖展開為二維形式并通過式（1）計算當前位置的貢獻值，與當前像素值相乘得到對應的顯著特征，如式（2）、（3）所示：

其中：xi、xj為像素值，Q(x)為當前像素經過Softmax 函數(shù)后的所得值，M(x)為經過特征提煉后的新特征值。通過特征提煉，對特征貢獻大的像素點信息保留，對貢獻小的像素點剔除置0 處理，得到新的特征矩陣MF，最后將新特征矩陣與師生網(wǎng)絡輸出的原始特征圖相乘，可得顯著特征矩陣，如式（4）、（5）所示：

其中：Ft、Fs為教師學生網(wǎng)絡的原始輸出矩陣，、為新的顯著特征矩陣。特征提煉模塊如圖2 所示。

2.3 距離計算模塊

2.3.1 類間距離計算模塊

語義分割是基于像素的分類任務，相比于教師網(wǎng)絡而言，學生網(wǎng)絡的結構相對簡單，參數(shù)少，對應的識別精度較低，類間距離也會相應的較小。針對特征提煉后的特征圖，本文通過對應坐標來編碼類間類別距離，類別總數(shù)可表示為：

其中：N代表圖像語義分割的類別總數(shù)；ci表示第i類的相對位置，其是通過特征提煉所得到的平均值。設Di，j為第i類和第j類類別中心的歐氏距離，如式（7）所示：

教師網(wǎng)絡的網(wǎng)絡深，參數(shù)量大。為了使學生網(wǎng)絡更好地在類間距離上模擬教師網(wǎng)絡，我們設計了類間距離損失函數(shù)Linter(T，S)，T為教師網(wǎng)絡，S為學生網(wǎng)絡，如式（8）所示：

2.3.2 類內距離計算模塊

不同于分類網(wǎng)絡，分割網(wǎng)絡往往是區(qū)域內的像素通過預測展現(xiàn)分類結果。除了類間距離，對于每個類來說，不同的像素預測值也是包含信息的。從教師蒸餾來的先驗信息占有更大比例的數(shù)值往往會“引領”對應的類別，而分配在該類中過小的數(shù)值有可能在下一輪預測中分配到其他類別中。因此，我們希望能夠提取到相應的類內距離信息從而提高分割的準確度。定義類內距離損失函數(shù)Lintra(T，S)如式（9）所示：

其中：E為第i類的邊緣像素特征點數(shù)量，ei表示第i類的邊緣特征點，ei0表示特征提煉得到的該類別對應特征點。距離計算模塊如圖3 所示。

圖3 距離計算模塊Fig.3 Distance calculation module

2.3.3 知識蒸餾損失函數(shù)設計

對于整體知識蒸餾損失函數(shù)設計，我們采用與文獻［8］相似的計算方式。該方法通過分別計算教師網(wǎng)絡和學生網(wǎng)絡對應的輸出特征圖的內積得到自己的相似矩陣，然后使用均方差計算兩個相似矩陣，使得教師網(wǎng)絡和學生網(wǎng)絡在同一類別與不同類別的分類上產生相似的結果輸出，同時結合類間距離和類內距離，使學生網(wǎng)絡在對應的特征映射上保留與教師網(wǎng)絡相似的知識。因此，模型的損失函數(shù)由類間損失函數(shù)Linter、類內距離損失函數(shù)Lintra以及與標簽的交叉熵損失函數(shù)Lcls3 部分組成，如式（10）所示：

其中：γ為平衡因子，ps為學生網(wǎng)絡的原始預測，g為真實標簽。整體算法偽代碼如算法1 所示。

3 實驗結果及分析

3.1 實驗設置及數(shù)據(jù)集介紹

為了驗證所提出的基于知識蒸餾的語義分割模型的有效性，我們在兩個公開數(shù)據(jù)集上進行實驗：（1）Cityscape［11］是一個包含19 類城市道路場景的數(shù)據(jù)集，包含了2 975 張訓練圖像、500 張驗證圖像以及1 525 張測試圖像。（2）Pascal VOC［12］由1 464 張訓練圖像、1 449 張驗證圖像以及1 456 張測試圖像組成，是一個用來做圖像分割與目標檢測的數(shù)據(jù)集。實驗環(huán)境選擇CPU 為AMD Ryzen 9 處理器，內存為64G，GPU 使 用RTX3090，采用梯度下降法訓練模型，優(yōu)化器選擇Adam 并動態(tài)調整學習率，初始學習率設置為0.002。另外，由于數(shù)據(jù)集分類數(shù)與復雜程度不同，我們設置了不同的平衡因子。對于VOC 數(shù)據(jù)集，γ設置為0.5；對于Cityscapes，γ設置為0.4。

為了更好地評估模型效果，我們選擇平均交并比（Mean Intersection over Union，MIoU）作為結果的評價標準。MIoU 如式（11）所示：

其中：pij表示真實值為i，被預測為j的數(shù)量；k+1是類別個數(shù)（包含空類別）；pii是真正的數(shù)量。MIoU 一般都是基于類別進行計算的，將每一類的IoU 計算后累加再進行平均得到基于全局的評價。

對于網(wǎng)絡架構，在本文所有的實驗中選擇DeepLabV3+［13］作為語義分割網(wǎng)絡的主體架構。經 過ImageNet［14］預訓練 的ResNet101［15］作為教師模型的骨干網(wǎng)絡。對于學生模型，我們使用ResNet18 作為骨干網(wǎng)絡。

3.2 消融實驗

所設計的蒸餾損失共包含3個損失函數(shù)：Linter、Lintra以及Lcls。為了驗證每一個損失函數(shù)的有效性，我們在Cityscapes 驗證集上進行了驗證實驗。教師網(wǎng)絡采用ResNet101 作為骨干網(wǎng)絡的Deep-LabV3+，學生網(wǎng)絡選擇ResNet18 作為骨干網(wǎng)絡的DeepLabV3+，實驗結果如表1 所示。其中，T代表教師網(wǎng)絡，S代表學生網(wǎng)絡。從表1 中可以看出，特征提煉模塊的添加使得分割精度提高至74.33%，類間距離損失函數(shù)Linter使模型分割精度由74.33%提高到75.61%，類內距離損失函數(shù)Lintra也使網(wǎng)絡分割精度提高了近0.7%。使用本文設計的蒸餾模型后，整體的分割精度達到76.53%，接近教師網(wǎng)絡77.76%的準確度。實驗證明我們提出的特征提煉模塊與蒸餾損失函數(shù)是有效的；（1）特征提煉模塊具有較好的分離特征與背景噪聲的能力，并對類內差異與類間距離有著較明顯的分割效果；（2）通過類內距離與類間距離分割獲得了更豐富的教師網(wǎng)絡知識，較基準網(wǎng)絡提高了6.17%。綜上所述，我們提出的方法提高了教師網(wǎng)絡對學生網(wǎng)絡的知識輸出，改善了師生網(wǎng)絡間知識的傳遞。

表1 消融實驗Tab.1 Ablation experiment

調整損失函數(shù)中的平衡因子γ是提高整體網(wǎng)絡性能的關鍵。在針對平衡因子的消融實驗中，在Pascal VOC 數(shù)據(jù)集 中，將γ設置為0.3、0.4、0.5 和0.6；在Cityscapes 數(shù)據(jù)集中，將γ設置為0.3、0.4、0.5。實驗結果如圖4 所示，可以看出，在Pascal VOC 數(shù)據(jù)集上的實驗中，γ=0.3 時，整個知識蒸餾的提升效果并不顯著；而當γ=0.6時，知識蒸餾的權重過高，自身預測結果對比真實標簽的損失函數(shù)影響過低，忽略了真實值對整體模型的監(jiān)督。Cityscapes 數(shù)據(jù)集的表現(xiàn)同樣如此，γ=0.4 時的效果最好。

圖4 超參數(shù)消融實驗結果Fig.4 Results of hyperparameter ablation experiment

3.3 對比實驗

3.3.1 針對Cityscapes 數(shù)據(jù)集的對比實驗

為了進一步驗證本文提出方法的有效性，我們與主流的蒸餾方法進行對比，如SKD（Knowledge adaptation for efficient semantic segmentation）［16］、SKDD（Structured knowledge distillation for dense prediction）［17］、IFVD［9］以及CD［10］，實驗在相同的網(wǎng)絡主架構上進行，實驗結果如表2 所示。由于不同的知識蒸餾方法均不改變原始原生網(wǎng)絡架構，故學生網(wǎng)絡對應不同方法的參數(shù)量與計算量均為14.07M 與270.5G。

表2 不同方法的對比實驗Tab.2 Comparison of experimental results of different methods

從表2 可以看出，學生網(wǎng)絡的分割結果與教師網(wǎng)絡相差甚微，其中模型參數(shù)量與每秒浮點運算次數(shù)（FLOPs）差別明顯，在知識蒸餾后學生網(wǎng)絡的分割效果得到了提高。與主流的知識蒸餾方法相比，我們所提出的方法的分割效果在原始的學生模型中有著顯著的提高，在Pascal VOC 與Cityscapes 數(shù)據(jù)集上分別有2.04%與4.45%的提升，由此能夠證明我們提出的方法在多個數(shù)據(jù)集中有著更好的分割性能。

我們在Cityscapes 測試集上對比了一些輕量化網(wǎng)絡模型，實驗結果如表3 所示。

表3 不同方法在Cityscapes 測試集上的對比實驗Tab.3 Comparative experiment of different methods in Cityscapes

本文設計的特征提煉模塊細化了內部特征，充分分離了前景特征與背景噪聲，不額外添加網(wǎng)絡計算量。與教師網(wǎng)絡相比，我們的參數(shù)量與FLOPs 大幅降低，分割精度僅降低了1.21%。與其他輕量化模型相比，我們所設計的模型有著顯著優(yōu)勢。類內距離計算模塊與類間距離計算模塊使得蒸餾傳遞的知識更加細化，特征內部特性更加清晰，經過知識蒸餾后的網(wǎng)絡模型分割精度比絕大部分輕量化模型都更高。與ICNet相比，我們的參數(shù)量只有其1/2，分割精度卻比ICNet 高7%；與RefineNet 相比，參數(shù)量大幅降低。學生網(wǎng)絡在分割精度相對較高的同時，參數(shù)量與FLOPs 都明顯更輕量化。

為了更直觀地對比圖像分割的性能差異，我們對原始學生網(wǎng)絡，CD［10］以及所提出的模型進行分割可視化，結果如圖5 所示。

圖5 Cityscapes 分割結果可視化Fig.5 Segmentation results of Cityscapes

根據(jù)分割結果可以看出，由于本文提出特征提煉的方法，分離了蒸餾知識的前景內容與背景噪聲，同時在特征空間的隱式編碼中提取類間距離和類內距離作為蒸餾損失，在圖像類別相近或內容重疊的部分相比于CD 分割的效果更好，如Cityscapes 中第三行的路邊欄桿所示，有著更清楚的劃分。

另外，我們還在Cityscapes 驗證集上對比了不同類別的MIoU，實驗結果如圖6 所示。

從圖6 中可以看到，所設計的方法對于每個小類別有著更好的分割表現(xiàn)。得益于對類間距離、類內距離知識蒸餾效果的提升，所設計的方法對一些重疊位置信息的內容有著更好的表現(xiàn)，如騎手、小汽車、自行車等。騎手的分割精度相比于IFVD 提升近7%，比CD提高了1.8%；自行車的分割精度比IFVD 提高了5%，比CD 提升了4.3%。

3.3.2 針對Pascal VOC 數(shù)據(jù)集的對比實驗

我們在Pascal VOC 驗證集上對比了不同網(wǎng)絡的參數(shù)和精度，如OCRNet（Object-contextual representations for semantic segmentation）［22］、DeeplabV3+［13］、FCN（Fully Convolutional Networks）［20］、ANN（Asymmetric non-local neural networks）［23］和PSPNet-ResNet101［24］。通過知識提煉以及類間與類內距離的蒸餾損失函數(shù)設計，DeeplabV3+-ResNet18 分別比FCN 和ANN 提高了7.81%和1.03%。實驗結果如圖7 所示。

圖7 Pascal VOC 驗證集分割結果對比Fig.7 Comparison of segmentation results of Pascal VOC validation

同時，我們采用ResNet18 與MobileNetV2 作為學生網(wǎng)絡在驗證集上驗證了所設計方法的有效性，實驗結果如圖8 所示。

圖8 不同骨干網(wǎng)絡對比Fig.8 Comparison of different backbone networks

當以MobileNetV2 作為學生網(wǎng)絡的骨干時，本文方法比知識蒸餾的原始網(wǎng)絡模型的分割準確率提高了4%，比IFVD 和CD 分別提高了2.1%和0.7%。在以ResNet18為骨干網(wǎng)絡的模型中，對比原始網(wǎng)絡，本文方法提高了2.03%；對比IFVD與CD，本文方法分別提高了1.52%與0.12%。我們同樣也對分割結果進行了可視化對比，如圖9所示。

圖9 Pascal VOC 分割結果可視化Fig.9 Segmentation results of Pascal VOC

可以看到，由于使用了類間距離與類內距離作為蒸餾損失，在空間位置相近的物體分割中，本文方法表現(xiàn)得更加優(yōu)秀，如第二行自行車與人體邊界的劃分，第三行中遠處建筑物的分割。

4 結 論

本文提出的基于知識蒸餾的特征提煉語義分割模型能夠更好地分離蒸餾知識的有效內容和噪聲，更準確地將特征信息傳遞給學生網(wǎng)絡。同時，本文網(wǎng)絡在特征空間的隱式編碼中提取類間距離與類內距離，從而獲取到相應的特征位置信息，之后學生網(wǎng)絡通過模擬特征位置信息來最小化與教師網(wǎng)絡特征位置的輸出，從而提高了學生網(wǎng)絡的分割精度。實驗結果表明，本文提出的方法在Pascal VOC 數(shù)據(jù)集上的MIoU 達到74.19%，在Cityscapes 數(shù)據(jù)集上的MIoU 達到76.53%，分別比原始學生網(wǎng)絡提高了2.04%和4.48%，具有較好的語義分割效果，證明了本文方法的有效性。