朱儁彥團隊提出GAN壓縮演算法：計算量減少20倍，生成效果不變，GPU、CPU統統能加速

原創 關注前沿科技 量子位

邊策 魚羊 發自 凹非寺

量子位 報道 | 公眾號 QbitAI

現如今，GAN的效果已經越來越出神入化。

比如英偉達的GauGAN，就如神筆馬良，能夠憑空造物：

不過，從無化有背後，計算量也相當驚人。

以GauGAN為例，與MobileNet-v3這樣的識別CNN相比，引數只相差1個數量級（4。2 ： 93）計算強度卻高出了2個數量級（0。5 ： 281）。

如此一來，互動式部署就變得很困難。

更直白來說，想要把模型部署到手機、平板這樣的終端，換裝變臉轉性別，這些邊緣裝置——吃不太消。

為了解決這一問題，來自MIT、Adobe研究院和上海交通大學的團隊琢磨出了一個通用壓縮框架。

有多強？

一舉將CycleGAN的計算量減少了20倍以上，將GauGAN的計算量減少了9倍，簡筆畫到實物圖的pix2pix也能駕馭，效果卻未差分毫。

值得一提的是。論文已入選CVPR 2020，程式碼也已開源。

作者團隊也星光璀璨，一作是來自上海交通大學ACM班的本科生李沐陽，CycleGAN的作者朱儁彥則親身參與、指導。

GAN壓縮演算法

壓縮GAN，並不像壓縮其他CNN模型那樣容易。

主要原因：GAN的訓練動力學高度不穩定，生成器與鑑別器之間存在巨大的結構差異，讓我們很難使用現有的CNN壓縮演算法。

為了解決這個問題，作者提出了針對有效生成模型量身定製的訓練方案，並透過神經架構搜尋（NAS）進一步提高壓縮率。

GAN壓縮的框架如下圖所示，主要分為3個部分：

1、給定一個預訓練的teacher生成器G’，透過蒸餾的方法獲取一個較小的“once-for-all”的student生成器G，其中包括透過權重共享的所有可能通道數。在每個訓練步驟中選擇不同的通道數ck。

2、從“once-for-all”生成器中提取許多子生成器，並評估其效能，無需重新訓練，這也是被叫做once-for-all（一勞永逸）的原因。

3、最後，根據給定的壓縮率目標和效能目標，選擇最佳子生成器，進行微調，並獲得最終的壓縮模型。

構造合適的損失函式

因為GAN壓縮演算法要面對CycleGAN、pix2pix還有GauGAN，這些模型的損失函式都不盡相同，所以需要構造一個新的損失函式。

統一未配對和配對學習

有些GAN是透過配對資料集學習的，有些則是非配對資料集。因此要在損失函式中加入第二項，統一非配對和配對學習的損失：

繼承teacher鑑別器

儘管目標是壓縮生成器，但是鑑別器會儲存學習到GAN知識，因為它會發現當前生成器的缺點。

因此，我們採用相同的鑑別器架構，使用teacher預訓練的權重，並與我們的壓縮生成器一起對鑑別器進行微調。預訓練的鑑別器可以指導student生成器的訓練。

初始化的student鑑別器D使用來自teacher鑑別器D’的權重。

中間特徵蒸餾

蒸餾是CNN中廣泛使用的模型壓縮方法。

CNN模型壓縮的一種廣泛使用的方法是知識蒸餾。透過匹配輸出層的logits，可以將知識從teacher模型轉移到student模型，從而提高後者的表現。

最後完整的損失函式為：

其中超引數λrecon和λdistill控制後兩項的權重。

高效的生成器設計空間

選擇設計良好的student體系結構對於最終進行知識蒸餾至關重要。

簡單地縮小teacher模型的通道數並不能產生緊湊的student模型，一旦計算量的壓縮比超過4倍，效能就會開始顯著下降。

預測可能的原因之一是，現有的生成器採用的影象識別模型，可能不是影象合成任務的最佳選擇。

下面，作者展示瞭如何從現有的cGAN生成器中獲得更好的架構設計空間，並在該空間內執行神經架構搜尋（NAS）。

卷積分解和層敏感性

近來高效的CNN設計，廣泛採用了卷積分解的方法，證明了在效能與計算之間的權衡取捨。作者發現使用分解的卷積，也有利於cGAN中的生成器設計。

然而實驗表明，將分解直接應用於所有卷積層，將大大降低影象質量。但是可以只有某些層分解會降低效能，而某些層則更魯棒。

在ResNet生成器中，resBlock層消耗了大部分模型引數和計算成本，而幾乎不受分解的影響。相反，上取樣層的引數要少得多，但是對模型壓縮相當敏感，適度的壓縮也會導致FID大幅下降。

NAS自動裁剪通道

現在的生成器在所有層上都使用手動設計，因此通道數會有冗餘。為了進一步壓縮模型，作者使用通道修剪自動選擇生成器的通道寬度減少冗餘，可以二次減少計算量。

對於每個卷積層，可以從8的倍數中選擇通道數，可以平衡MAC和硬體並行性。

解耦訓練和結構搜尋

為了解決該問題，作者遵循one-shot神經體系結構搜尋方法的最新工作，將模型訓練與體系結構搜尋分離 。

先訓練一個支援所有通道的“once-for-all”網路，具有不同數量通道的每個子網路都經過同等訓練，可以獨立執行，子網路與“once-for-all”網路共享權重。

在訓練了“once-for-all”網路後，透過直接在驗證集上評估每個候選子網路的效能來找到最佳子網。由於“once-for-all”網路經過權重共享的全面訓練，因此無需進行微調。

透過這種方式，我們可以將訓練和搜尋生成器體系結構分離開來：只需要訓練一次，在無需進一步訓練的情況下評估所有可能的通道配置，並選擇最佳的作為搜尋結果。

實驗結果

最終實驗結果如下：

GAN壓縮將乘法累加運算（MAC）的數量減了9~21倍。

這種方法將CycleGAN的計算量減少了20倍以上，將pix2pix的計算量減少了9倍，將GauGAN的計算量減少了9倍。

而且不僅能在GPU上加速，在各種各種各樣的CPU上也可以實現加速，而且加速效果在CPU上更明顯。

程式碼開源

現在，這一GAN壓縮框架已經開源，資料集和預訓練模型均可下載。

如果想要上手親測，需要的環境是：

Linux

Python 3

CPU or NVIDIA GPU + CUDA CuDNN

同時，研究團隊也準備了CycleGAN和pix2pix兩個版本的PyTorch Colab，線上即可體驗。

試試用CycleGAN把普通馬變斑馬的效果。

首先把GAN Compression這個專案克隆下來：

！git clone https：//github。com/mit-han-lab/gan-compression。git

下載測試資料集：

！bash datasets/download_cyclegan_dataset。sh horse2zebra

下載預訓練模型：

python scripts/download_model。py —model pix2pix —task edges2shoes-r —stage full

python scripts/download_model。py —model pix2pix —task edges2shoes-r —stage compressed

下面，就是見證換皮的時刻：

效果差距大不大，肉眼看了不算，還是要計算一下FID（用於評估GAN生成影象的質量，分數越低代表與真實影象越相似）。

該專案提供了幾個資料集的真實統計資訊：

bash 。/datasets/download_real_stat。sh horse2zebra A

bash 。/datasets/download_real_stat。sh horse2zebra B

測試的結果顯示，原始CycleGAN的FID是65。687，壓縮後，模型的FID是65。312，兩者差距不大。

但在MAC、引數規模和延遲方面，壓縮模型都要遠遠小於原始模型。

上海交大本科生一作，朱儁彥&韓松加持

論文一作，是上海交通大學ACM班大四本科生李沐陽。

2019年7月-今年1月，李沐陽師從MIT韓松教授和朱儁彥教授，在MIT Han Lab擔任研究助理，這篇CVPR 2020論文就是在此期間產出。

目前，李沐陽的研究興趣在於機器學習、系統以及計算機圖形學等領域，他計劃在今年畢業之後繼續讀博深造。

科研之餘，李沐陽同學的一大愛好是唱歌。在他的個人主頁上，他還分享了自己的唱吧連結，以及個人MV。

論文的另外幾位作者，分別是：

Ji Lin，本科畢業於清華大學，現於MIT就讀博士。

丁堯堯，同樣是來自上海交大ACM班的本科生，和李沐陽同在MIT Han Lab擔任研究助理，受韓松教授指導。

Zhijian Liu，本科畢業於上海交大，現於韓松教授門下就讀博士。

朱儁彥，李沐陽的指導者之一。這位青年大牛無需多介紹，CycleGAN作者，國際頂會ACM SIGGRAPH 2018最佳博士論文獎獲得者。現為Adobe研究科學家，今年秋天將回歸母校CMU擔任助理教授。

韓松，MIT EECS助理教授，同樣是AI業界大牛。博士畢業於斯坦福大學，曾斬獲ICLR 2016最佳論文、FPGA 2017最佳論文。

傳送門

論文地址：

https：//arxiv。org/abs/2003。08936

GitHub：

https：//github。com/mit-han-lab/gan-compression

原標題：《朱儁彥團隊提出GAN壓縮演算法：計算量減少20倍，生成效果不變，GPU、CPU統統能加速》