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