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