選自arXiv
作者：Muyang Li等
機器之心編譯
機器之心編輯部
GAN 也可以大幅壓縮，MIT 韓松團隊的最新研究讓眾多研究者們為之一振。
生成模型 GAN 是機器學習領域裡最為重要的發展方向之一。但這類算法需要消耗巨量算力，大多數研究者已經很難得出新成果。近年來，這一方向頗有被大型機構壟斷的趨勢。
但近日，來自麻省理工學院（MIT）、Adobe、上海交通大學的研究者提出了一種用於壓縮條件 GAN 的通用方法。這一新技術在保持視覺保真度的同時，將 pix2pix，CycleGAN 和 GauGAN 等廣泛使用的條件 GAN 模型的計算量減少到 1/9~1/21。該方法適用於多種生成器架構、學習目標，配對或非配對設置。
目前該研究的論文已被 CVPR 2020 大會收錄，有關 GAN 壓縮框架的 PyTorch 版實現也已開源。
項目連結：https://github.com/mit-han-lab/gan-compression
GAN 壓縮後的具體性能如何？在研究者們展示的 Demo 中，使用 CycleGAN 為視頻中的馬添加斑馬條紋所需的算力不到 1/16，幀數提高了三倍，而且效果還有所提升：
值得一提的是，該研究所使用的硬體平臺是英偉達的邊緣 AI 計算晶片 Jetson Xavier GPU。根據官方公布的數據，Jetson Xavier 的 INT8 算力為 22+10TOPS，驍龍 865 則是 15TOPS。壓縮後的 GAN 現在看來已經可以跑在機器人、無人機等小型設備上了，未來塞進手機指日可待。
論文連結：https://arxiv.org/pdf/2003.08936v1.pdf
總體介紹
生產對抗網絡（GANS）擅長合成十分逼真的圖像。GAN 的一種變體——條件式生成對抗網絡（conditional generative adversarial network，cGAN）在很多計算機視覺和圖像學應用中都能夠實現可控制的圖像合成。但這些應用大都需要模型與人互動，因此需要低延遲的設備上才能獲得更好的用戶體驗。
然而，近來推出的一些 cGAN 在計算強度上比當前識別卷積神經網絡（CNN）大 1 至 2 個量級。舉例而言，GanGAN 每張圖像消耗 281G Macs，而 MobileNet-v3 僅需 0.44G Macs，這就導致前者很難用於交互式部署。
而且，就目前來看，邊緣設備又多數受到內容以及電池之類硬體層面上的限制，也阻礙了 GAN 在邊緣設備上的部署。
因此，基於 GAN 和 cGAN 在圖像合成領域存在的這些問題，韓松團隊提出了 GAN 壓縮，這是一種通用壓縮方法，用於減少 GAN 的推理時間以及計算成本。同時，壓縮生成模型面臨兩個基本困難：GAN 訓練不穩定，尤其是在未配對的情況下；生成器與 CNN 不同，因此很難使用現有的 CNN 設計。為了解決此問題，團隊將知識從原始的教師生成器中間表示層傳遞到其相應的學生生成器層中。
為了降低訓練成本，團隊還通過訓練包含所有可能通道數的「once- for-all network」，將模型訓練與架構搜索分離。這個「once-for-all network」可以通過權重共享產生許多子網絡，無需訓練就可以評估每個子網絡的性能。該模型可應用至各種條件下的 GAN 模型，不管其屬於哪種模型架構、學習算法或監督設置（配對或未配對）。
通過大量的實驗，團隊已證明了此方法可以將 pix2pix，CycleGAN 以及 GauGAN 三種廣泛使用的 GAN 模型計算量減少至 1/9 到 1/21，同時還不會損失生成圖像的保真度。
具體方法
我們都知道，對用於交互式應用的條件式生成模型進行壓縮具有挑戰性，這主要是由以下兩方面原因造成的。其一，從本質上講，GAN 的動態訓練非常不穩定；其二，識別和生成模型之間存在的巨大架構差異導致很難直接使用現有的 CNN 壓縮算法。
基於這些原因，研究者提出了專門針對高效生成模型的訓練方案，並利用神經架構搜索（NAS）進一步增加壓縮比。GAN 壓縮框架整體架構如下圖 3 所示，其中他們利用 ResNet 生成器作為示例。需要強調的是，同一框架可以應用於不同的生成器架構和學習目標。
圖 3：文中 GAN 壓縮框架的整體架構圖。
目標函數
1. 統一配對學習和非配對學習
廣泛的訓練目標使得構建通用壓縮框架非常困難。為了解決這一問題，研究者在模型壓縮設置中統一了配對和非配對學習，不管教師模型最初是如何訓練的。給定原始教師生成器 G′，研究者將非配對訓練設置轉換為配對設置。對於非配對設置，可以將原始生成器輸出視為真值並使用配對訓練目標訓練壓縮後的生成器 G。
學習目標總結如下：
基於這些修改，現在可以將同一個壓縮框架應用於不同類型的 cGAN。而且，與原始的非配對訓練設置相比，使用上述偽對（pseudo pair）進行學習可以使訓練更加穩定，並產生更好的結果。
2. 從教師判別器中學習
儘管此研究致力於壓縮生成器，但判別器 D 中儲存著 GAN 的有用信息。因此，研究者利用了同一個判別器架構，使用來自教師判別器的預訓練權重，與壓縮生成器一起微調判別器。
在實驗中，研究者觀察到，預訓練判別器可以指導學生生成器的訓練。使用一個隨機初始化的判別器通常會導致訓練不穩定及圖像質量下降。這個 GAN 的目標可以寫成以下形式：
在上式中，研究者使用來自教師判別器 D′的權重來初始化學生判別器 D。他們使用一個標準的極小化極大優化器訓練 G 和 D。
3. 中間特徵蒸餾
知識蒸餾是 CNN 模型壓縮用到的一種普遍方法。通過匹配輸出層 logit 的分布，可以將來自教師模型的暗知識（dark knowledge）遷移到學生模型中，提高學生模型的性能。然而，條件 GAN 通常會輸出一個確定的圖像，而不是概率分布。
為了解決上述問題，研究者匹配了教師生成器的中間表示。中間層包含的通道越多，它們所能提供的信息越豐富，學生模型所獲取的輸出之外的信息也越多。蒸餾目標如下：
其中，G_t(x) 和 G′_t(x) 是學生和教師模型中第 t 個被選層的中間特徵激活，T 表示層數。
4. 完整優化目標
最終目標可以寫成如下形式：
其中，超參數λ_recon 和 λ_distill 控制每個項的重要性。
高效生成器設計空間
選擇一個設計良好的學生架構對最終知識蒸餾的效果是至關重要的，研究者發現，光是縮減教師模型的通道數量並不能使學生模型更緊湊：當計算縮減量超過 4 倍時，性能就會顯著下降。
1. 卷積分解和層靈敏度
現有的生成器通常採用傳統卷積來匹配 CNN 分類和分段設計。近來一些高效的 CNN 設計廣泛採用卷積的分解形式（depthwise + pointwise），在性能和計算二者之間的均衡性更好。研究者發現，分解的卷積也可以用在 cGAN 的生成器設計上。
2. 使用 NAS 實現自動裁剪通道
現有的生成器在所有層上使用手動設計（並且幾乎統一）的通道數，如此一來就會產生冗餘，遠非最優方法。為了進一步提升壓縮效率，研究者使用通道剪枝（channel pruning）來自動選擇生成器中的通道寬度，從而減少冗餘，二次減少計算量。這一方法支持有關通道數的細粒度選擇，針對每個卷積層，可以從 8 的倍數中選擇卷積層，從而平衡 MAC 和硬體並行性。
解耦訓練與結構搜索
研究者依照最近 one-shot 的 NAS 方法的研究，將模型訓練與架構搜索脫鉤。首先，訓練一個支持不同通道數量的「once-for-all」網絡，其中的每個子網絡都經過了同等訓練。圖 3 闡釋了整個框架。研究者假設原始的教師生成器有
個通道，對於給定的通道數
，從「once-for-all」的權重張量提取第一個
通道，獲得子網絡的權重網絡。
在每一個訓練步驟中，使用學習目標對具備某一通道數的子網絡進行隨機取樣，計算輸出和梯度，更新提取的權重（公式 4）。由於最先被抽取的幾個通道更新頻率更高，它們在所有權重之中扮演著更為關鍵的角色。
這個「once-for-all」網絡訓練好之後，研究者直接在驗證集上評估了每個子網絡的性能，找出了最佳的子網絡。「once-for-all」網絡已經經過了權值共享的徹底訓練，無需再進行微調（fine-tuning）。這個結果近似於從頭開始訓練模型的性能。
通過這種方式，只需要訓練一次，且無需再進一步訓練，就能評估所有評估所有通道的配置，並依據搜索結果找到其中最佳。當然，也可以對選中的架構進行微調來進一步提升其性能。
實驗結果
研究者在以下三種條件式 GAN 模型上進行實驗以驗證文中 GAN 壓縮框架的泛化性，它們分別是 CycleGAN、Pix2Pix 和 GauGAN。所使用的四個數據集為 Horsezebra、Edgesshoes、Cityscapes 和 Mapaerial photo。
下表 1 展示了在上述四個數據集上壓縮 CycleGAN、Pix2Pix 和 GauGAN 模型的量化結果。
表 1：三種條件式 GAN 模型壓縮的量化評估，其中 Cityscapes 數據集上使用 mAP 度量（越高越好），其他數據集上使用 FID 度量。結果顯示，在性能輕微下降的情況下，文中 GAN 壓縮方法可以在 MACs 上將當前 SOTA 條件式 GAN 壓縮 7 至 21 倍，在模型大小上壓縮 5 至 33 倍。對於 CycleGAN 模型壓縮，文中 GAN 壓縮方法遠優於以往的 CycleGAN-specific Co-evolution 方法。
性能與計算之間的權衡
該方法除了能夠實現較大的壓縮率以外，同時其也可以提升不同模型大小的性能。下圖 6 顯示了在 pix2pix 模型中，在不同數據集上的性能與計算的權衡。
圖 6：在 Cityscapes 與 EdgesShoes 數據集上 pix2pix 的權衡曲線。剪枝與蒸餾方法在大型模型中超出了從頭開始訓練，然而在模型被急劇壓縮時表現較差。
效果展示
下圖 4 為使用該方法得到的效果展示。圖中分別給出了輸入數據、標準輸出、原始模型輸出以及壓縮後模型的輸出。從圖中可以看出，即使在較大的壓縮率下，研究者所提方法仍然能保持輸出圖像的視覺可信度。
圖 4：Cityscapes、EdgesShoes 以及 HorseZebra 數據集上效果比較。
硬體推理加速
對於真實場景交互應用而言，在硬體設備上推理加速的重要性要遠大於減少計算消耗。如下表 2 所示，為驗證所提方法在實際應用中的有效性，研究者在具有不同運算性能的設備上面測試了壓縮模型的推理速度。
表 2：在 NVIDIA Jetson AGX Xavier、NVIDIA Jetson Nano、1080Ti GPU 和 Xeon CPU 上測到的內存和延遲下降。
結論
在本文中，韓松團隊提出的通用壓縮框架能夠顯著降低條件式 GAN 中生成器的計算花銷和模型大小，並且通過知識蒸餾和自然架構搜索來提升訓練穩定性以及模型效率。實驗表明，文中提出的 GAN 壓縮方法能夠在保持視覺質量的同時壓縮數種條件式 GAN 模型。他們表示，未來的研究工作將側重於降低模型延遲以及構建生成視頻模型的高效框架。
首期「AI 開發者成長計劃」由機器之心聯合微眾銀行開設《聯邦學習 FATE 入門與應用實戰》公開課，並特邀合作夥伴 VMware 帶來為期 4 周共 6 期課程，設置主題講解、項目實操、在線答疑等環節，幫助 AI 開發者從零入門聯邦學習。公開課第四課將於 3 月 24 日開講，歡迎大家報名參與。