朱俊彥團(tuán)隊(duì)提出GAN壓縮算法：計(jì)算量減少20倍，生成效果不變，GPU、CPU統(tǒng)統(tǒng)能加速

邊策 魚羊 發(fā)自 凹非寺
量子位 報(bào)道 | 公眾號 QbitAI

現(xiàn)如今，GAN的效果已經(jīng)越來越出神入化。

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

不過，從無化有背后，計(jì)算量也相當(dāng)驚人。

以GauGAN為例，與MobileNet-v3這樣的識別CNN相比，參數(shù)只相差1個數(shù)量級（4.2 : 93）計(jì)算強(qiáng)度卻高出了2個數(shù)量級（0.5 : 281）。

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

更直白來說，想要把模型部署到手機(jī)、平板這樣的終端，換裝變臉轉(zhuǎn)性別，這些邊緣設(shè)備——吃不太消。

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

有多強(qiáng)？

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

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

作者團(tuán)隊(duì)也星光璀璨，一作是來自上海交通大學(xué)ACM班的本科生李沐陽，CycleGAN的作者朱俊彥則親身參與、指導(dǎo)。

GAN壓縮算法

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

主要原因：GAN的訓(xùn)練動力學(xué)高度不穩(wěn)定，生成器與鑒別器之間存在巨大的結(jié)構(gòu)差異，讓我們很難使用現(xiàn)有的CNN壓縮算法。

為了解決這個問題，作者提出了針對有效生成模型量身定制的訓(xùn)練方案，并通過神經(jīng)架構(gòu)搜索（NAS）進(jìn)一步提高壓縮率。

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

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

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

3、最后，根據(jù)給定的壓縮率目標(biāo)和性能目標(biāo)，選擇最佳子生成器，進(jìn)行微調(diào)，并獲得最終的壓縮模型。

構(gòu)造合適的損失函數(shù)

因?yàn)镚AN壓縮算法要面對CycleGAN、pix2pix還有GauGAN，這些模型的損失函數(shù)都不盡相同，所以需要構(gòu)造一個新的損失函數(shù)。

統(tǒng)一未配對和配對學(xué)習(xí)

有些GAN是通過配對數(shù)據(jù)集學(xué)習(xí)的，有些則是非配對數(shù)據(jù)集。因此要在損失函數(shù)中加入第二項(xiàng)，統(tǒng)一非配對和配對學(xué)習(xí)的損失：

繼承teacher鑒別器

盡管目標(biāo)是壓縮生成器，但是鑒別器會存儲學(xué)習(xí)到GAN知識，因?yàn)樗鼤l(fā)現(xiàn)當(dāng)前生成器的缺點(diǎn)。

因此，我們采用相同的鑒別器架構(gòu)，使用teacher預(yù)訓(xùn)練的權(quán)重，并與我們的壓縮生成器一起對鑒別器進(jìn)行微調(diào)。預(yù)訓(xùn)練的鑒別器可以指導(dǎo)student生成器的訓(xùn)練。

初始化的student鑒別器D使用來自teacher鑒別器D’的權(quán)重。

中間特征蒸餾

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

CNN模型壓縮的一種廣泛使用的方法是知識蒸餾。通過匹配輸出層的logits，可以將知識從teacher模型轉(zhuǎn)移到student模型，從而提高后者的表現(xiàn)。

最后完整的損失函數(shù)為：

其中超參數(shù)λrecon和λdistill控制后兩項(xiàng)的權(quán)重。

高效的生成器設(shè)計(jì)空間

選擇設(shè)計(jì)良好的student體系結(jié)構(gòu)對于最終進(jìn)行知識蒸餾至關(guān)重要。

簡單地縮小teacher模型的通道數(shù)并不能產(chǎn)生緊湊的student模型，一旦計(jì)算量的壓縮比超過4倍，性能就會開始顯著下降。

預(yù)測可能的原因之一是，現(xiàn)有的生成器采用的圖像識別模型，可能不是圖像合成任務(wù)的最佳選擇。

下面，作者展示了如何從現(xiàn)有的cGAN生成器中獲得更好的架構(gòu)設(shè)計(jì)空間，并在該空間內(nèi)執(zhí)行神經(jīng)架構(gòu)搜索（NAS）。

卷積分解和層敏感性

近來高效的CNN設(shè)計(jì)，廣泛采用了卷積分解的方法，證明了在性能與計(jì)算之間的權(quán)衡取舍。作者發(fā)現(xiàn)使用分解的卷積，也有利于cGAN中的生成器設(shè)計(jì)。

然而實(shí)驗(yàn)表明，將分解直接應(yīng)用于所有卷積層，將大大降低圖像質(zhì)量。但是可以只有某些層分解會降低性能，而某些層則更魯棒。

在ResNet生成器中，resBlock層消耗了大部分模型參數(shù)和計(jì)算成本，而幾乎不受分解的影響。相反，上采樣層的參數(shù)要少得多，但是對模型壓縮相當(dāng)敏感，適度的壓縮也會導(dǎo)致FID大幅下降。

NAS自動裁剪通道

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

對于每個卷積層，可以從8的倍數(shù)中選擇通道數(shù)，可以平衡MAC和硬件并行性。

解耦訓(xùn)練和結(jié)構(gòu)搜索

為了解決該問題，作者遵循one-shot神經(jīng)體系結(jié)構(gòu)搜索方法的最新工作，將模型訓(xùn)練與體系結(jié)構(gòu)搜索分離 。

先訓(xùn)練一個支持所有通道的“once-for-all”網(wǎng)絡(luò)，具有不同數(shù)量通道的每個子網(wǎng)絡(luò)都經(jīng)過同等訓(xùn)練，可以獨(dú)立運(yùn)行，子網(wǎng)絡(luò)與“once-for-all”網(wǎng)絡(luò)共享權(quán)重。

在訓(xùn)練了“once-for-all”網(wǎng)絡(luò)后，通過直接在驗(yàn)證集上評估每個候選子網(wǎng)絡(luò)的性能來找到最佳子網(wǎng)。由于“once-for-all”網(wǎng)絡(luò)經(jīng)過權(quán)重共享的全面訓(xùn)練，因此無需進(jìn)行微調(diào)。

通過這種方式，我們可以將訓(xùn)練和搜索生成器體系結(jié)構(gòu)分離開來：只需要訓(xùn)練一次，在無需進(jìn)一步訓(xùn)練的情況下評估所有可能的通道配置，并選擇最佳的作為搜索結(jié)果。

實(shí)驗(yàn)結(jié)果

最終實(shí)驗(yàn)結(jié)果如下：

GAN壓縮將乘法累加運(yùn)算（MAC）的數(shù)量減了9~21倍。

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

而且不僅能在GPU上加速，在各種各種各樣的CPU上也可以實(shí)現(xiàn)加速，而且加速效果在CPU上更明顯。

代碼開源

現(xiàn)在，這一GAN壓縮框架已經(jīng)開源，數(shù)據(jù)集和預(yù)訓(xùn)練模型均可下載。

如果想要上手親測，需要的環(huán)境是：

• Linux
• Python 3
• CPU or NVIDIA GPU + CUDA CuDNN

同時，研究團(tuán)隊(duì)也準(zhǔn)備了CycleGAN和pix2pix兩個版本的PyTorch Colab，在線即可體驗(yàn)。

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

首先把GAN Compression這個項(xiàng)目克隆下來：

!git?clone?https://github.com/mit-han-lab/gan-compression.git

下載測試數(shù)據(jù)集：

!bash?datasets/download_cyclegan_dataset.sh?horse2zebra

下載預(yù)訓(xùn)練模型：

python?scripts/download_model.py?—model?pix2pix?—task?edges2shoes-r?—stage?full
python?scripts/download_model.py?—model?pix2pix?—task?edges2shoes-r?—stage?compressed

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

效果差距大不大，肉眼看了不算，還是要計(jì)算一下FID（用于評估GAN生成圖像的質(zhì)量，分?jǐn)?shù)越低代表與真實(shí)圖像越相似）。

該項(xiàng)目提供了幾個數(shù)據(jù)集的真實(shí)統(tǒng)計(jì)信息：

bash?./datasets/download_real_stat.sh?horse2zebra?A
bash?./datasets/download_real_stat.sh?horse2zebra?B

測試的結(jié)果顯示，原始CycleGAN的FID是65.687，壓縮后，模型的FID是65.312，兩者差距不大。

但在MAC、參數(shù)規(guī)模和延遲方面，壓縮模型都要遠(yuǎn)遠(yuǎn)小于原始模型。

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

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

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

目前，李沐陽的研究興趣在于機(jī)器學(xué)習(xí)、系統(tǒng)以及計(jì)算機(jī)圖形學(xué)等領(lǐng)域，他計(jì)劃在今年畢業(yè)之后繼續(xù)讀博深造。

科研之余，李沐陽同學(xué)的一大愛好是唱歌。在他的個人主頁上，他還分享了自己的唱吧鏈接，以及個人MV。

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

Ji Lin，本科畢業(yè)于清華大學(xué)，現(xiàn)于MIT就讀博士。

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

Zhijian Liu，本科畢業(yè)于上海交大，現(xiàn)于韓松教授門下就讀博士。

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

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

傳送門

論文地址：https://arxiv.org/abs/2003.08936

GitHub：https://github.com/mit-han-lab/gan-compression

李沐陽個人主頁：https://lmxyy.me

Colab：https://colab.research.google.com/github/mit-han-lab/gan-compression/blob/master/cycle_gan.ipynbhttps://colab.research.google.com/github/mit-han-lab/gan-compression/blob/master/pix2pix.ipynb

— 完 —