知識蒸餾for語義分割
Structured Knowledge Distillation for Semantic Segmentation
2019/03/13 Author:Yu ZHang
2019cvpr中的一篇文章,是我見到的第一個在分割上使用知識蒸餾的,可見實時的分割已經是現在研究的熱門了,作者做的事情就是:用PSPNet或OCNet等作為老師網路,來指導學生網路,如最近看的ESPNet,MobileNet,ShuffleNet等,讓這些小的網路擁有更高的準確率。下圖是整個蒸餾的過程:
整個蒸餾分為三個部分:1)Pixle-wise distillation 2)Pair-wise distillation 3)Holistic distillation
- Pair-wise distillation
首先教師網路是不進行優化的,其是已經訓好的且好使的分割網路。按先後順序來說我們先講Pair-wise distillation,這一個部分作者是受馬爾科夫隨機場(條件隨機場)的啟發,作者想要找兩兩畫素之間的相關性以提高網路的效果。At表示教師網路特徵圖第i個畫素與第j個畫素之間的相關性,As表示學生網路ij之間相關性,通過下式(平方差)來計算蒸餾loss,讓學生網路逼近教師網路:
而計算兩點之間相似性公式如下:
- Pixel-wise distillation
這個蒸餾模組就是正常蒸餾的思路,我們把教師網路輸出的概率拿過來,與學生網路輸出的概率做loss,讓學生網路逼近教師網路,在圖中直觀的看就是兩個分割圖做loss,但其實是概率做loss,loss公式如下:
- Holistic distillation
這裡作者利用了GAN的思想,學生網路被看作是為生成器,其輸入就是資料集中的RGB影象,輸出的是分割圖(fake),而教師網路輸出的分割圖為真實圖,分割圖與RGB圖送入Discriminator net做一個嵌入,這個網路相當於GAN中的判別器會輸出一個得分,得分表示RGB與分割圖是否匹配,真實分佈與錯誤分佈之間計算Wasserstein距離。
最終共有4個loss(還有一個loss是學生網路輸出與真實label做普通的loss),最終loss公式如下:
採用此蒸餾方法,作者將許多小網路的準確率都提升了非常多,實驗結果如下圖: