[論文レビュー] MobileStyleGAN: A Lightweight Convolutional Neural Network for High-Fidelity Image Synthesis

Sergei Belousov|arXiv (Cornell University)|Apr 10, 2021

Generative Adversarial Networks and Image Synthesis参考文献 27被引用数 13

ひとこと要約

MobileStyleGANは、StyleGAN2と比較してパラメータを3.5倍、計算複雑性を9.5倍削減する軽量でウェーブレットベースのスタイルベースGANアーキテクチャを提案する。1024×1024のFFHQデータセットでFID=7.75（StyleGAN2は2.84）という同等の画像品質を達成する。事前学習済みのStyleGAN2教師モデルからの知識蒸留、深度分離可能モジュレート畳み込みの導入、グラフ最適化とOpenVINOによる加速により、エッジデバイスでの効率的な推論を可能にする。

ABSTRACT

In recent years, the use of Generative Adversarial Networks (GANs) has become very popular in generative image modeling. While style-based GAN architectures yield state-of-the-art results in high-fidelity image synthesis, computationally, they are highly complex. In our work, we focus on the performance optimization of style-based generative models. We analyze the most computationally hard parts of StyleGAN2, and propose changes in the generator network to make it possible to deploy style-based generative networks in the edge devices. We introduce MobileStyleGAN architecture, which has x3.5 fewer parameters and is x9.5 less computationally complex than StyleGAN2, while providing comparable quality.

研究の動機と目的

• 最新のスタイルベースGAN（例：StyleGAN2）の高い計算コストがエッジデバイスへのデプロイを制限する問題に対処すること。
• 高解像度画像合成において、画像品質を損なわずにモデルサイズと推論複雑性を低減すること。
• アーキテクチャの再設計とモデル圧縮技術を通じて、CPUおよびモバイルプラットフォームでの効率的な推論を可能にすること。
• 効率性と滑らかな潜在空間の両立を図るため、生成モデルにウェーブレット表現を統合する可能性を検討すること。
• 大規模な教師ネットワーク（StyleGAN2）からコンactな学生ネットワーク（MobileStyleGAN）へ高精細画像生成能力を転移するための蒸留パイプラインを開発すること。

提案手法

• 出力画像の離散ウェーブレット変換（DWT）を予測するウェーブレットベースの画像表現を採用することで、周波数ドメイン処理が可能になる。
• 標準的なモジュレート畳み込みの代替として、FLOPsを削減しつつスタイルモジュレーションを維持する軽量な深さ分離可能モジュレート畳み込み（DSMC）を導入する。
• グラフレベル最適化（例：演算統合、定数畳み込み）と互換性を持つ、改訂されたデモジュレーション機構を提案する。
• 学生ネットワークが教師（StyleGAN2）を、知覚損失、GAN損失、ピクセルレベル損失を用いて模倣する知識蒸留によりMobileStyleGANを訓練する。
• L1ピクセル損失、VGG16ベースの知覚損失、R1正則化付きの敵対的GAN損失を組み合わせたマルチ損失目的関数を用い、学習の安定化を図る。
• 微分可能オーグメンテーション（アフィン変換、カットアウト）とOpenVINOを活用し、CPU上で最適化された推論を実現し、リアルタイムな顔生成を可能にする。

実験結果

リサーチクエスチョン

• RQ1 GANにおけるウェーブレットベースの周波数ドメイン表現は、高精細画像合成を維持しつつ、計算複雑性を顕著に低減できるか？
• RQ2 深度分離可能モジュレート畳み込みは、画像品質に悪影響を及げることなく、標準的なモジュレート畳み込みにどの程度置き換え可能か？
• RQ3 知識蒸留は、大規模なStyleGAN2モデルからコンパクトなMobileStyleGAN学生モデルへ高精細画像生成能力を効果的に転送できるか？
• RQ4 アーキテクチャの変更により、高価なハードウェアに依存せずにCPUやエッジデバイスでの効率的推論を実現できるか？
• RQ5 ウェーブレット変換と最適化された正規化の統合は、潜在空間の滑らかさと学習の安定性を向上させるか？

主な発見

• MobileStyleGANは、StyleGAN2の28.27Mから8.01Mへとパラメータを削減し、モデルサイズを3.5倍に縮小した。
• 計算複雑性は143.15 GMACsから15.09 GMACsに低下し、FLOPsで9.5倍の削減を達成した。
• FFHQデータセットにおいて、MobileStyleGANのFréchet Inception Distance（FID）は7.75であり、StyleGAN2の2.84と比較して、圧縮にもかかわらず高い知覚的品質を示した。
• CPU（Intel i5-8279U）上での推論時間は、PyTorch使用時で4.3秒からOpenVINO使用時で0.16秒にまで短縮され、強力な最適化の可能性が示された。
• 定数畳み込みや演算統合などのグラフ最適化を活用したOpenVINOの使用により、リアルタイム推論が可能となり、MobileStyleGANはエッジデプロイに適したものとなった。
• 蒸留パイプラインは、知覚損失と敵対的損失により、高精細でリアルな、シャープな出力を効果的に転送できた。