GAN 也可以大幅壓縮，MIT 韓松團(tuán)隊(duì)的最新研究讓眾多研究者們?yōu)橹徽瘛?/p>

[[319638]]

生成模型 GAN 是機(jī)器學(xué)習(xí)領(lǐng)域里最為重要的發(fā)展方向之一。但這類算法需要消耗巨量算力，大多數(shù)研究者已經(jīng)很難得出新成果。近年來(lái)，這一方向頗有被大型機(jī)構(gòu)壟斷的趨勢(shì)。

但近日，來(lái)自麻省理工學(xué)院（MIT）、Adobe、上海交通大學(xué)的研究者提出了一種用于壓縮條件 GAN 的通用方法。這一新技術(shù)在保持視覺(jué)保真度的同時(shí)，將 pix2pix，CycleGAN 和 GauGAN 等廣泛使用的條件 GAN 模型的計(jì)算量減少到 1/9~1/21。該方法適用于多種生成器架構(gòu)、學(xué)習(xí)目標(biāo)，配對(duì)或非配對(duì)設(shè)置。

目前該研究的論文已被 CVPR 2020 大會(huì)收錄，有關(guān) GAN 壓縮框架的 PyTorch 版實(shí)現(xiàn)也已開源。

項(xiàng)目鏈接：https://github.com/mit-han-lab/gan-compression

GAN 壓縮后的具體性能如何？在研究者們展示的 Demo 中，使用 CycleGAN 為視頻中的馬添加斑馬條紋所需的算力不到 1/16，幀數(shù)提高了三倍，而且效果還有所提升：

值得一提的是，該研究所使用的硬件平臺(tái)是英偉達(dá)的邊緣 AI 計(jì)算芯片 Jetson Xavier GPU。根據(jù)官方公布的數(shù)據(jù)，Jetson Xavier 的 INT8 算力為 22+10TOPS，驍龍 865 則是 15TOPS。壓縮后的 GAN 現(xiàn)在看來(lái)已經(jīng)可以跑在機(jī)器人、無(wú)人機(jī)等小型設(shè)備上了，未來(lái)塞進(jìn)手機(jī)指日可待。

論文鏈接：https://arxiv.org/pdf/2003.08936v1.pdf

總體介紹

生產(chǎn)對(duì)抗網(wǎng)絡(luò)（GANS）擅長(zhǎng)合成十分逼真的圖像。GAN 的一種變體——條件式生成對(duì)抗網(wǎng)絡(luò)（conditional generative adversarial network，cGAN）在很多計(jì)算機(jī)視覺(jué)和圖像學(xué)應(yīng)用中都能夠?qū)崿F(xiàn)可控制的圖像合成。但這些應(yīng)用大都需要模型與人互動(dòng)，因此需要低延遲的設(shè)備上才能獲得更好的用戶體驗(yàn)。

然而，近來(lái)推出的一些 cGAN 在計(jì)算強(qiáng)度上比當(dāng)前識(shí)別卷積神經(jīng)網(wǎng)絡(luò)（CNN）大 1 至 2 個(gè)量級(jí)。舉例而言，GanGAN 每張圖像消耗 281G Macs，而 MobileNet-v3 僅需 0.44G Macs，這就導(dǎo)致前者很難用于交互式部署。

而且，就目前來(lái)看，邊緣設(shè)備又多數(shù)受到內(nèi)容以及電池之類硬件層面上的限制，也阻礙了 GAN 在邊緣設(shè)備上的部署。

因此，基于 GAN 和 cGAN 在圖像合成領(lǐng)域存在的這些問(wèn)題，韓松團(tuán)隊(duì)提出了 GAN 壓縮，這是一種通用壓縮方法，用于減少 GAN 的推理時(shí)間以及計(jì)算成本。同時(shí)，壓縮生成模型面臨兩個(gè)基本困難：GAN 訓(xùn)練不穩(wěn)定，尤其是在未配對(duì)的情況下；生成器與 CNN 不同，因此很難使用現(xiàn)有的 CNN 設(shè)計(jì)。為了解決此問(wèn)題，團(tuán)隊(duì)將知識(shí)從原始的教師生成器中間表示層傳遞到其相應(yīng)的學(xué)生生成器層中。

為了降低訓(xùn)練成本，團(tuán)隊(duì)還通過(guò)訓(xùn)練包含所有可能通道數(shù)的「once- for-all network」，將模型訓(xùn)練與架構(gòu)搜索分離。這個(gè)「once-for-all network」可以通過(guò)權(quán)重共享產(chǎn)生許多子網(wǎng)絡(luò)，無(wú)需訓(xùn)練就可以評(píng)估每個(gè)子網(wǎng)絡(luò)的性能。該模型可應(yīng)用至各種條件下的 GAN 模型，不管其屬于哪種模型架構(gòu)、學(xué)習(xí)算法或監(jiān)督設(shè)置（配對(duì)或未配對(duì)）。

通過(guò)大量的實(shí)驗(yàn)，團(tuán)隊(duì)已證明了此方法可以將 pix2pix，CycleGAN 以及 GauGAN 三種廣泛使用的 GAN 模型計(jì)算量減少至 1/9 到 1/21，同時(shí)還不會(huì)損失生成圖像的保真度。

具體方法

我們都知道，對(duì)用于交互式應(yīng)用的條件式生成模型進(jìn)行壓縮具有挑戰(zhàn)性，這主要是由以下兩方面原因造成的。其一，從本質(zhì)上講，GAN 的動(dòng)態(tài)訓(xùn)練非常不穩(wěn)定；其二，識(shí)別和生成模型之間存在的巨大架構(gòu)差異導(dǎo)致很難直接使用現(xiàn)有的 CNN 壓縮算法。

基于這些原因，研究者提出了專門針對(duì)高效生成模型的訓(xùn)練方案，并利用神經(jīng)架構(gòu)搜索（NAS）進(jìn)一步增加壓縮比。GAN 壓縮框架整體架構(gòu)如下圖 3 所示，其中他們利用 ResNet 生成器作為示例。需要強(qiáng)調(diào)的是，同一框架可以應(yīng)用于不同的生成器架構(gòu)和學(xué)習(xí)目標(biāo)。

圖 3：文中 GAN 壓縮框架的整體架構(gòu)圖。

目標(biāo)函數(shù)

1. 統(tǒng)一配對(duì)學(xué)習(xí)和非配對(duì)學(xué)習(xí)

廣泛的訓(xùn)練目標(biāo)使得構(gòu)建通用壓縮框架非常困難。為了解決這一問(wèn)題，研究者在模型壓縮設(shè)置中統(tǒng)一了配對(duì)和非配對(duì)學(xué)習(xí)，不管教師模型最初是如何訓(xùn)練的。給定原始教師生成器 G′，研究者將非配對(duì)訓(xùn)練設(shè)置轉(zhuǎn)換為配對(duì)設(shè)置。對(duì)于非配對(duì)設(shè)置，可以將原始生成器輸出視為真值并使用配對(duì)訓(xùn)練目標(biāo)訓(xùn)練壓縮后的生成器 G。

學(xué)習(xí)目標(biāo)總結(jié)如下：

基于這些修改，現(xiàn)在可以將同一個(gè)壓縮框架應(yīng)用于不同類型的 cGAN。而且，與原始的非配對(duì)訓(xùn)練設(shè)置相比，使用上述偽對(duì)（pseudo pair）進(jìn)行學(xué)習(xí)可以使訓(xùn)練更加穩(wěn)定，并產(chǎn)生更好的結(jié)果。

2. 從教師判別器中學(xué)習(xí)

盡管此研究致力于壓縮生成器，但判別器 D 中儲(chǔ)存著 GAN 的有用信息。因此，研究者利用了同一個(gè)判別器架構(gòu)，使用來(lái)自教師判別器的預(yù)訓(xùn)練權(quán)重，與壓縮生成器一起微調(diào)判別器。

在實(shí)驗(yàn)中，研究者觀察到，預(yù)訓(xùn)練判別器可以指導(dǎo)學(xué)生生成器的訓(xùn)練。使用一個(gè)隨機(jī)初始化的判別器通常會(huì)導(dǎo)致訓(xùn)練不穩(wěn)定及圖像質(zhì)量下降。這個(gè) GAN 的目標(biāo)可以寫成以下形式：

在上式中，研究者使用來(lái)自教師判別器 D′的權(quán)重來(lái)初始化學(xué)生判別器 D。他們使用一個(gè)標(biāo)準(zhǔn)的極小化極大優(yōu)化器訓(xùn)練 G 和 D。

3. 中間特征蒸餾

知識(shí)蒸餾是 CNN 模型壓縮用到的一種普遍方法。通過(guò)匹配輸出層 logit 的分布，可以將來(lái)自教師模型的暗知識(shí)（dark knowledge）遷移到學(xué)生模型中，提高學(xué)生模型的性能。然而，條件 GAN 通常會(huì)輸出一個(gè)確定的圖像，而不是概率分布。

為了解決上述問(wèn)題，研究者匹配了教師生成器的中間表示。中間層包含的通道越多，它們所能提供的信息越豐富，學(xué)生模型所獲取的輸出之外的信息也越多。蒸餾目標(biāo)如下：

其中，G_t(x) 和 G′_t(x) 是學(xué)生和教師模型中第 t 個(gè)被選層的中間特征激活，T 表示層數(shù)。

4. 完整優(yōu)化目標(biāo)

最終目標(biāo)可以寫成如下形式：

其中，超參數(shù)λ_recon 和 λ_distill 控制每個(gè)項(xiàng)的重要性。

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

選擇一個(gè)設(shè)計(jì)良好的學(xué)生架構(gòu)對(duì)最終知識(shí)蒸餾的效果是至關(guān)重要的，研究者發(fā)現(xiàn)，光是縮減教師模型的通道數(shù)量并不能使學(xué)生模型更緊湊：當(dāng)計(jì)算縮減量超過(guò) 4 倍時(shí)，性能就會(huì)顯著下降。

1. 卷積分解和層靈敏度

現(xiàn)有的生成器通常采用傳統(tǒng)卷積來(lái)匹配 CNN 分類和分段設(shè)計(jì)。近來(lái)一些高效的 CNN 設(shè)計(jì)廣泛采用卷積的分解形式（depthwise + pointwise），在性能和計(jì)算二者之間的均衡性更好。研究者發(fā)現(xiàn)，分解的卷積也可以用在 cGAN 的生成器設(shè)計(jì)上。

2. 使用 NAS 實(shí)現(xiàn)自動(dòng)裁剪通道

現(xiàn)有的生成器在所有層上使用手動(dòng)設(shè)計(jì)（并且?guī)缀踅y(tǒng)一）的通道數(shù)，如此一來(lái)就會(huì)產(chǎn)生冗余，遠(yuǎn)非最優(yōu)方法。為了進(jìn)一步提升壓縮效率，研究者使用通道剪枝（channel pruning）來(lái)自動(dòng)選擇生成器中的通道寬度，從而減少冗余，二次減少計(jì)算量。這一方法支持有關(guān)通道數(shù)的細(xì)粒度選擇，針對(duì)每個(gè)卷積層，可以從 8 的倍數(shù)中選擇卷積層，從而平衡 MAC 和硬件并行性。

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

研究者依照最近 one-shot 的 NAS 方法的研究，將模型訓(xùn)練與架構(gòu)搜索脫鉤。首先，訓(xùn)練一個(gè)支持不同通道數(shù)量的「once-for-all」網(wǎng)絡(luò)，其中的每個(gè)子網(wǎng)絡(luò)都經(jīng)過(guò)了同等訓(xùn)練。圖 3 闡釋了整個(gè)框架。研究者假設(shè)原始的教師生成器有個(gè)通道，對(duì)于給定的通道數(shù)

，從「once-for-all」的權(quán)重張量提取第一個(gè)通道，獲得子網(wǎng)絡(luò)的權(quán)重網(wǎng)絡(luò)。

在每一個(gè)訓(xùn)練步驟中，使用學(xué)習(xí)目標(biāo)對(duì)具備某一通道數(shù)的子網(wǎng)絡(luò)進(jìn)行隨機(jī)取樣，計(jì)算輸出和梯度，更新提取的權(quán)重（公式 4）。由于最先被抽取的幾個(gè)通道更新頻率更高，它們?cè)谒袡?quán)重之中扮演著更為關(guān)鍵的角色。

這個(gè)「once-for-all」網(wǎng)絡(luò)訓(xùn)練好之后，研究者直接在驗(yàn)證集上評(píng)估了每個(gè)子網(wǎng)絡(luò)的性能，找出了最佳的子網(wǎng)絡(luò)?！竜nce-for-all」網(wǎng)絡(luò)已經(jīng)經(jīng)過(guò)了權(quán)值共享的徹底訓(xùn)練，無(wú)需再進(jìn)行微調(diào)（fine-tuning）。這個(gè)結(jié)果近似于從頭開始訓(xùn)練模型的性能。

通過(guò)這種方式，只需要訓(xùn)練一次，且無(wú)需再進(jìn)一步訓(xùn)練，就能評(píng)估所有評(píng)估所有通道的配置，并依據(jù)搜索結(jié)果找到其中最佳。當(dāng)然，也可以對(duì)選中的架構(gòu)進(jìn)行微調(diào)來(lái)進(jìn)一步提升其性能。

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

研究者在以下三種條件式 GAN 模型上進(jìn)行實(shí)驗(yàn)以驗(yàn)證文中 GAN 壓縮框架的泛化性，它們分別是 CycleGAN、Pix2Pix 和 GauGAN。所使用的四個(gè)數(shù)據(jù)集為 Horse↔zebra、Edges→shoes、Cityscapes 和 Map↔aerial photo。

下表 1 展示了在上述四個(gè)數(shù)據(jù)集上壓縮 CycleGAN、Pix2Pix 和 GauGAN 模型的量化結(jié)果。

表 1：三種條件式 GAN 模型壓縮的量化評(píng)估，其中 Cityscapes 數(shù)據(jù)集上使用 mAP 度量（越高越好），其他數(shù)據(jù)集上使用 FID 度量。結(jié)果顯示，在性能輕微下降的情況下，文中 GAN 壓縮方法可以在 MACs 上將當(dāng)前 SOTA 條件式 GAN 壓縮 7 至 21 倍，在模型大小上壓縮 5 至 33 倍。對(duì)于 CycleGAN 模型壓縮，文中 GAN 壓縮方法遠(yuǎn)優(yōu)于以往的 CycleGAN-specific Co-evolution 方法。

性能與計(jì)算之間的權(quán)衡

該方法除了能夠?qū)崿F(xiàn)較大的壓縮率以外，同時(shí)其也可以提升不同模型大小的性能。下圖 6 顯示了在 pix2pix 模型中，在不同數(shù)據(jù)集上的性能與計(jì)算的權(quán)衡。

圖 6：在 Cityscapes 與 Edges→Shoes 數(shù)據(jù)集上 pix2pix 的權(quán)衡曲線。剪枝與蒸餾方法在大型模型中超出了從頭開始訓(xùn)練，然而在模型被急劇壓縮時(shí)表現(xiàn)較差。

效果展示

下圖 4 為使用該方法得到的效果展示。圖中分別給出了輸入數(shù)據(jù)、標(biāo)準(zhǔn)輸出、原始模型輸出以及壓縮后模型的輸出。從圖中可以看出，即使在較大的壓縮率下，研究者所提方法仍然能保持輸出圖像的視覺(jué)可信度。

圖 4：Cityscapes、Edges→Shoes 以及 Horse→Zebra 數(shù)據(jù)集上效果比較。

硬件推理加速

對(duì)于真實(shí)場(chǎng)景交互應(yīng)用而言，在硬件設(shè)備上推理加速的重要性要遠(yuǎn)大于減少計(jì)算消耗。如下表 2 所示，為驗(yàn)證所提方法在實(shí)際應(yīng)用中的有效性，研究者在具有不同運(yùn)算性能的設(shè)備上面測(cè)試了壓縮模型的推理速度。

表 2：在 NVIDIA Jetson AGX Xavier、NVIDIA Jetson Nano、1080Ti GPU 和 Xeon CPU 上測(cè)到的內(nèi)存和延遲下降。

結(jié)論

在本文中，韓松團(tuán)隊(duì)提出的通用壓縮框架能夠顯著降低條件式 GAN 中生成器的計(jì)算花銷和模型大小，并且通過(guò)知識(shí)蒸餾和自然架構(gòu)搜索來(lái)提升訓(xùn)練穩(wěn)定性以及模型效率。實(shí)驗(yàn)表明，文中提出的 GAN 壓縮方法能夠在保持視覺(jué)質(zhì)量的同時(shí)壓縮數(shù)種條件式 GAN 模型。他們表示，未來(lái)的研究工作將側(cè)重于降低模型延遲以及構(gòu)建生成視頻模型的高效框架。