IJCAI21 對比自蒸餾實現(xiàn)超分模型的加速

背景 Background

卷積神經(jīng)網(wǎng)絡(luò)在超分任務(wù)上取得了很好的成果，但是依然存在著參數(shù)繁重、顯存占用大、計算量大的問題，為了解決這些問題，作者提出利用對比自蒸餾實現(xiàn)超分模型的壓縮和加速。

主要貢獻

• 作者提出的對比自蒸餾（CSD）框架可以作為一種通用的方法來同時壓縮和加速超分網(wǎng)絡(luò)，在落地應(yīng)用中的運行時間也十分友好。
• 自蒸餾被引用進超分領(lǐng)域來實現(xiàn)模型的加速和壓縮，同時作者提出利用對比學(xué)習(xí)進行有效的知識遷移，從而 進一步的提高學(xué)生網(wǎng)絡(luò)的模型性能。
• 在Urban100數(shù)據(jù)集上，加速后的EDSR+可以實現(xiàn)4倍的壓縮比例和1.77倍的速度提高，帶來的性能損失僅為0.13 dB PSNR。

方法 ?Method

CSSR-Net

作者提出的CSD框架主要包含兩個部分：CSSR-Net和Contrastive loss。先簡單介紹一下 channel-splitting super-resolution network (CSSR-Net), CSSR-Net其實可以是任何超分網(wǎng)絡(luò)的子網(wǎng)絡(luò)，與原網(wǎng)絡(luò)的區(qū)別僅在于channel數(shù)量上，在CSD框架中作為學(xué)生網(wǎng)絡(luò)，來作為加速優(yōu)化的目標產(chǎn)物模型。

如圖中所示，在CSD框架中，CSSR-Net被耦合進教師網(wǎng)絡(luò)，共享教師網(wǎng)絡(luò)的部分參數(shù)（即共同參數(shù)部分）。

在訓(xùn)練過程中聯(lián)合優(yōu)化的目標函數(shù)為：

作者發(fā)現(xiàn)，如果只有單獨的知識蒸餾（即不加對比損失來進一步提高模型性能），訓(xùn)練后得到的模型效果其實更差，詳見表1。

作者推測其原因為：隱式的知識蒸餾不夠強，無法通過兩項獨立的損失項來提供足夠insightful的信息。因此作者引進對比學(xué)習(xí)來顯示的在學(xué)生和教師之間建立一種聯(lián)系,為優(yōu)化目標提供了一個封閉的上下限，從而同時提高學(xué)生網(wǎng)絡(luò)和教師網(wǎng)絡(luò)的性能。上限被構(gòu)建來將CSSR-Net的輸出向教師拉近，下限來限制CSSR-Net的輸出遠離負樣本。

Contrastive Learning

論文中的對比學(xué)習(xí)損失其實非常簡單，基本和論文作者團隊在AECR-Net（“Contrastive learning for single image dehazing”）中的contrastive loss一致：

其中，是VGG預(yù)訓(xùn)練網(wǎng)絡(luò)的第層，M是總隱藏層的數(shù)量，是損失，我們希望上式的分母越大越好，分子越小越好。這種損失引入了相反的力，將CSSR-Net的輸出向其教師的輸出特征拉近，向負樣本的特征拉遠。

感知損失與CL相比，其只有一個優(yōu)化的上限來限制學(xué)生網(wǎng)絡(luò)的輸出，與感知損失不同，Contrastive learning使用了多種負樣本來作為一個優(yōu)化的下限目標減小解空間，并且進一步提高CSSR—Net和教師網(wǎng)絡(luò)的性能。

實驗

• 如原文圖6所示，引入對比學(xué)習(xí)損失計算的負樣本越多，模型性能越好，這很好解釋，更多的負樣本進一步約束了下限，每個batch的loss來自于更好的數(shù)據(jù)采樣。

• 表2中值得注意的一點：沒有負樣本約束的 CSD（GT Pos），相比與baseline來說，性能更差，作者推測 這是由于HR圖像只是提供了更強的上限，我們的模型容量是有限的，也就是說模型其實是欠擬合的，這種更強的上限反而難以完全釋放模型的性能。（We speculate that HR images provide a stronger upper bound which is more difficult for the limited capacity S to fully exlpoit. ）

筆者觀點

這篇論文還是很有趣的，通過自蒸餾來實現(xiàn)超分模型的加速和壓縮，出發(fā)點很好，但是性能上的提高感覺不是非常顯著，對比學(xué)習(xí)損失其實和作者團隊cvpr21的對比學(xué)習(xí)去霧基本一致。這篇論文讓我想到cvpr20中一篇利用memory unit實現(xiàn)自蒸餾去雨的論文，也很有趣。