邊策 魚羊 發自 凹非寺
量子位 報道 | 公衆号 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
李沐陽個人首頁:
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
—完—
@量子位 · 追蹤AI技術和産品新動态
深有感觸的朋友,歡迎贊同、關注、分享三連վ'ᴗ' ի ❤