微軟亞研院提出用於語義分割的結構化知識蒸餾

科技 03-16

作者 | CV君

來源 | 我愛計算機視覺

今天跟大家分享一篇關於語義分割的論文，剛剛上傳到arXiv的CVPR 2019接收論文《Structured Knowledge Distillation for Semantic Segmentation》，通訊作者單位為微軟亞洲研究院。

作者信息：

作者分別來自澳大利亞阿德萊德大學、微軟亞洲研究院、北航、Keep公司、三星中國研究院，該文為第一作者Yifan Liu在微軟亞洲研究院實習期間的工作。

該文研究了在語義分割模型的知識蒸餾中引入結構化信息的損失函數，在不改變模型計算量的情況下，使用該方法在Cityscapes數據集上mIoU精度取得了最高達15.17%的提升。

什麼是知識蒸餾？

顧名思義，知識蒸餾是把知識濃縮到「小」網路模型中。一般情況下，在相同的數據上訓練，模型參數量較大、計算量大的模型往往精度比較高，而用精度高、模型複雜度高的模型即Teacher網路的輸出訓練Student網路，以期達到使計算量小參數少的小網路精度提升的方法，就是知識蒸餾。

知識蒸餾的好處是顯而易見的，使用知識蒸餾後的Student網路能夠達到較高的精度，而且更有利於實際應用部署，尤其是在移動設備中。

下面兩幅圖中，作者展示了使用該文提出的結構化知識蒸餾的語義分割模型在計算量和參數量不變的情況下，精度獲得了大幅提升。

演算法原理

知識蒸餾的目標是希望對於Teacher網路和Student網路給定相同的圖像，輸出結果盡量一樣。

所以，知識蒸餾的關鍵，是如何衡量Teacher網路和Student網路輸出結果的一致性，也就是訓練過程中的損失函數設計。

該文中作者將語義分割問題看為像素分類問題，所以很自然的可以使用衡量分類差異的逐像素（Pixel-wise）的損失函數Cross entropy loss，這是在最終的輸出結果Score map中計算的。

同時作者引入了圖像的結構化信息損失，如下圖所示。

如何理解圖像的結構化信息？一種很顯然的結構化信息即圖像中局部的一致性。在語義分割中，可以簡單理解為，預測結果中存在的自相似性，作者衡量這種結構化信息的方式是Teacher預測的兩像素結果和Student網路預測的兩像素結果一致。衡量這種損失，作者稱之為Pair-wise loss（也許可以翻譯為「逐成對像素」損失）。

另一種更高層次的結構化信息是來自對圖像整體結構相似性的度量，作者引入了對抗網路的思想，設計專門的網路分支分類Teacher網路和Student網路預測的結果，網路收斂的結果是該網路不能再區分Teacher網路和Student網路的輸出。作者稱這塊損失函數為Holistic loss（整體損失）。

仔細想想，作者設計的損失函數的三部分，逐像素的損失（Pixel-wise loss，PI）、逐像素對的損失（Pair-wise loss，PA）、整體損失（Holistic loss，HO）都很有道理，是不是？

作者使用ResNet18網路模型在Cityscapes數據集上研究了作者提出的損失函數各部分對結果的影響。（ImN代表用ImageNet預訓練模型初始化網路）

結果如下圖。