壓縮版styleGAN，合成高保真圖像，參數(shù)更少、計(jì)算復(fù)雜度更低
近年來在生成圖像建模中，生成對抗網(wǎng)絡(luò)（GAN）的應(yīng)用越來越多?；跇邮剑╯tyle-based）的 GAN 可以生成不同層次的細(xì)節(jié)，大到頭部形狀、小到眼睛顏色，它在高保真圖像合成方面實(shí)現(xiàn)了 SOTA，但其生成過程的計(jì)算復(fù)雜度卻非常高，難以應(yīng)用于智能手機(jī)等移動設(shè)備。
近日，一項(xiàng)專注于基于樣式的生成模型的性能優(yōu)化的研究引發(fā)了大家的關(guān)注。該研究分析了 StyleGAN2 中最困難的計(jì)算部分，并對生成器網(wǎng)絡(luò)提出了更改，使得在邊緣設(shè)備中部署基于樣式的生成網(wǎng)絡(luò)成為可能。該研究提出了一種名為 MobileStyleGAN 的新架構(gòu)。相比于 StyleGAN2，該架構(gòu)的參數(shù)量減少了約 71 %，計(jì)算復(fù)雜度降低約 90 %，并且生成質(zhì)量幾乎沒有下降。
StyleGAN2（上）與 MobileStyleGAN（下）的生成效果對比。
論文作者已將 MobileStyleGAN 的 PyTorch 實(shí)現(xiàn)放到了 GitHub 上。
論文地址：
https://arxiv.org/pdf/2104.04767.pdf
項(xiàng)目地址：
https://github.com/bes-dev/MobileStyleGAN.pytorch
該實(shí)現(xiàn)所需的訓(xùn)練代碼非常簡單：
StyleGAN2（左）與 MobileStyleGAN（右）的生成效果展示。
下面我們來具體看一下 MobileStyleGAN 架構(gòu)的方法細(xì)節(jié)。
MobileStyleGAN 架構(gòu)
MobileStyleGAN 架構(gòu)是在基于樣式生成模型的基礎(chǔ)上構(gòu)建的，它包括映射網(wǎng)絡(luò)和合成網(wǎng)絡(luò)，前者采用的是 StyleGAN2 中的映射網(wǎng)絡(luò)，該研究的重點(diǎn)是設(shè)計(jì)了一個計(jì)算高效的合成網(wǎng)絡(luò)。
MobileStyleGAN 與 StyleGAN2 的區(qū)別
StyleGAN2 使用基于像素的圖像表征，并旨在直接預(yù)測輸出圖像的像素值。而 MobileStyleGAN 使用基于頻率的圖像表征，旨在預(yù)測輸出圖像的離散小波變換 (DWT)。當(dāng)應(yīng)用到 2D 圖像，DWT 將信道轉(zhuǎn)換成四個大小相同的信道，這幾個信道具有較低的空間分辨率和不同的頻帶。然后，逆向離散小波變換(IDWT) 從小波域重建基于像素的表征，如下圖所示。
StyleGAN2 利用跳遠(yuǎn)生成器（skip-generator），通過對同一圖像的多個分辨率的 RGB 值進(jìn)行顯式求和來形成輸出圖像。該研究發(fā)現(xiàn)，當(dāng)在小波域中對圖像進(jìn)行預(yù)測時(shí)，基于跳遠(yuǎn)連接（skip connection）的預(yù)測頭對生成圖像的質(zhì)量影響不大。因此，為了降低計(jì)算復(fù)雜度，該研究采用網(wǎng)絡(luò)中最后一個塊的單個預(yù)測頭替換跳遠(yuǎn)生成器。但從中間塊中預(yù)測目標(biāo)圖像對于穩(wěn)定的圖像合成具有重要意義。因此，該研究為每個中間塊添加一個輔助預(yù)測頭，根據(jù)目標(biāo)圖像的空間分辨率對其進(jìn)行預(yù)測。
StyleGAN2 和 MobileStyleGAN 的預(yù)測頭區(qū)別。
如下圖所示，調(diào)制卷積包括調(diào)制、卷積和歸一化（左）。深度可分離調(diào)制卷積也包括這些部分（中）。StyleGAN2 描述了用于權(quán)重的調(diào)制 / 解調(diào)，該研究分別將它們應(yīng)用于輸入 / 輸出激活，這使得描述深度可分離調(diào)制卷積更加容易。
StyleGAN2 構(gòu)造塊使用 ConvTranspose（下圖左）來 upscale 輸入特征映射。而該研究在 MobileStyleGAN 構(gòu)造塊（下圖右）中使用 IDWT 當(dāng)作 upscale 函數(shù)。由于 IDWT 不包含可訓(xùn)練參數(shù)，該研究在 IDWT 層之后增加了額外的深度可分離調(diào)制卷積。
StyleGAN2 和 MobileStyleGAN 的完整構(gòu)造塊結(jié)構(gòu)如下圖所示：
基于蒸餾的訓(xùn)練過程
類似于此前的一些研究，該研究的訓(xùn)練框架也基于知識蒸餾技術(shù)。該研究將 StyleGAN2 作為教師網(wǎng)絡(luò)，訓(xùn)練 MobileStyleGAN 來模仿 StyleGAN2 的功能，訓(xùn)練框架如下圖所示。
