たった一行のコードでStyleGAN2の高画質画像を生成する【NNabla】

AI実装

StyleGAN2は本物と見分けのつかないような高画質な画像を生成できるAIアルゴリズムです。ですが、実装面ではTensorFlow1.X系が必要でありGPU:RTX3000シリーズ※の方は非常に使いずらい仕様となっています。Dockerを使って回避する方法がありますが、実装難易度が高く抵抗ある方も多いと思います。

そんな方々に朗報です。NNablaを使うことで簡単にStyleGAN2の画像生成を試すことができます!今回はその方法を説明していきます。NNablaはソニーのニューラルネットワークライブラリーで現在の主要なディープラーニングアルゴリズムがサポートされています。実装も簡単にできるので、かなりオススメです。

Neural Network Libraries by Sony
Neural Network Libraries by Sony is the open source software to make research, development and implementation of neural network more efficient.

ちなみに今回はローカル環境での実装を目指しますが、GoogleColabで良いよーという方はnnabla examples Github中段あたりのStyleGAN2の”Open in Colab”をクリックするとそのまま実行できます。

※RTX3000シリーズはCUDA10系をサポートしていないので、TensorFlow1.Xを使えません😅

StyleGAN2のアーキテクチャについて

StyleGAN2はかなり複雑なアルゴリズムになっています。StyleGAN2を説明するためには、まずStyleGANの説明から必要になります。実装を急ぐ方はここは読み飛ばしてください。

https://arxiv.org/abs/1812.04948

StyleGANは画像の解像度を段階的に上げていくProgressive Growing技術を用いています。↑図のSynthesis network gの灰色の枠右下に4×4や8×8と記載されています。これは解像度で最初は4x4pixelからスタートし徐々に解像度を上げていく仕組みになっています。

また、畳み込み層の間にAdaIN層が挿入されています。このAdaINによって各畳み込み層へ画像のStyleを取り込む形になっています。AdaINへは潜在変数zをMapping Networkという全結合ニューラルネットワークを通し非線形変換したstyleベクトルWを入力しています。これにより潜在変数zを変更することでStyleを変えた画像を生成でき、StyleGANという名前になっています。

StyleGAN2はStyleGANの弱点を対策したモデルとなっています。StyleGANの弱点は↓

  • 水滴のようなノイズが生成されることがある
  • 画像の一部がStyle変更に追随せず、頻出する特徴を生成してしまう

StyleGAN2で実装している内容は↓

  • AdaINで実際の統計量を正規化するのではなく、推定の統計量を正規化することで不要な情報は顕在化してしまうのを防ぐ=>水滴ノイズを防ぐ
  • Progressive Growingに変え、ResNetのようなskip connectionをもつGeneratorへ変更し頻出する特徴に引っ張られてしまう現象を対策
  • 画像が潜在空間上で滑らかさの指標Perceptual Path Length(PPL)を正規化項としてモデル組み込み画質向上

StyleGAN2のアルゴリズムは上述のようにかなり複雑です。より詳細の内容は↓のサイトに分かりやすくまとまっていました。ご参考に。

GANの基礎からStyleGAN2まで
この記事について

NNablaでStyleGAN2の画像生成をやってみる

NNablaを使ってローカル環境でStyleGAN2の画像生成を試してみます。上述のようにStyleGAN2の実装はTensorFlow1.X系が必要でしたが、NNablaをつかうことでTensorFlowなしで実装することができます。マシンスペックに不安の方はGoogleClabからも実行できます。このページをスクロールしていき、StyleGAN2の”Open in Colab”をクリックすると実装することができます。

使用PCスペック

  • Windows 10
  • CPU : AMD Ryzen 7 5800
  • GPU : RTX3070 8GB

事前準備としてCUDA、cudnnのインストールします。この辺りはここに詳しく掲載されているので、その通りにインストールしています。インストールしたバージョンは下記です。

  • NVIDIA CUDA Toolkit 11.4
  • NVIDIA cuDNN v8.2.0

ソースコード

nnabla examplesをgit clone かダウンロードしてください。必要なライブラリは下記です。AnacondaPromptなどでpipインストールしてください。

  • nnabla
  • python >= 3.6

事前学習データのダウンロード

NNabla example/stylegan2Pretrained Weightの表にあるdownloadより事前学習データをダウンロードしてください。事前学習データは「人の顔」「車」「猫」「馬」「教会」「ImageNet」の6種類サポートされており、お好みで選んでOKです。ダウンロードしたデータを先ほどダウンロードしたソースコードnnabla-examples/GANs/stylegan2ディレクトリ下へ保存してください。

コード実行

NNablaをつかうと本当に簡単にStyleGAN2の画像生成をすることができます。なんとコードはこの一行です!↓のコードを実行するだけで基本OKです。AnacondaPromptなどで実行してみてください。

python generate.py --seed 217 --batch-size 4 --dataset-type human_faces

引数を簡単に説明すると、

  • –seed : 潜在変数zの値です。ここへ入れる数値で出力画像が変化します。
  • –batch-size : 生成する画像の枚数です。
  • –dataset-type : 生成する画像の種類です。上記でダウンロードした事前学習データに合わせて変更することで様々な画像を生成できます。

resultsディレクトリを確認してみてください。こんな感じの1024x1024pixelの高画質な顔画像が生成できたと思います。これが実在しないAIがつくり出した人物というので驚きですよね。

ほかの引数についても試してみます。

python generate.py --seed 217  --truncation-psi 0.3  --stochastic-seed 1993
  • –truncation-psi : Wが平均値に近くなるように調整します。数値が小さいと画質が向上する一方でバリエーションが失われます。
  • –stochastic-seed : Noise B(上記StyleGANのブロック図参照)の値を指定。肌質やそばかすなど見た目に大きな影響を与えない特徴をコントロールします。
–truncation-psi 0.5 –stochastic-seed 1
–truncation-psi 0.3 –stochastic-seed 1993

引数を変更することで、帽子がなくなって少し成長した感じになりました。こんな感じで調整できるところもStyleGAN2の面白いところですね。

Style Mixingを試してみる

StyleGANにはStyle Mixingという面白い技術があります。Synthesis networkに入力されるスタイルをある層から他の画像のベクトルに入れ替えることによって、2つの画像の特徴をミックス出来る技術です。難しいことは置いといて、↓の画像で分かると思います。画像下段の女性の髪、肌の色が、上段の画像より引き継がれています。これがStyleMixingです。

style_mixing_sample.png
https://github.com/sony/nnabla-examples/tree/master/GANs/stylegan2

折角なので猫の画像でやってみようと思います。まずは猫の画像を2種類生成してみます。↓のコードを入れて実行してください。

python generate.py --seed 200 --batch-size 1 --dataset-type cat --truncation-psi 0.7
python generate.py --seed 500 --batch-size 1 --dataset-type cat --truncation-psi 0.7
–seed 200
–seed 500

猫の画像が生成されたと思います。これも本物と見分けのつかないですよね。この2つの猫をかけ合わせて画像生成してみます。↓のコードを実行してください。

python generate.py --seed 200 --mixing --seed-mix 500 --mix-after 6 --batch-size 1 --dataset-type cat --truncation-psi 0.7
Style Mixing sample

こんな感じで上の二つの猫がミックスされた猫が生成されました。引数をすこし説明すると↓です。

  • –mixing : Style Mixingを適用する場合このコードを記入する。
  • –seed-mix : seedで指定した潜在変数で生成された画像と、seed-mixで指定した潜在変数で生成された画像をStyle Mixingします。
  • –mix-after : どの層からStyle Mixingするかを決定します。値が小さいとseed-mixの特長を色濃く反映するようになります。

StyleMixingで遊んでみます。↓図の下段左端の猫と上段の猫の画像(StyleGAN2で生成した画像)をStyle Mixingしてみます。下段の猫をベースとして、毛色や瞳孔が上段の猫とStyle Mixingできていると思います。癒されますね~🐈こんなことが簡単に実装できるので是非お試しあれ。

最後にStyleGAN2がつくりだした面白画像をいくつか紹介します。StyleGAN2も完璧では無いんですね。

マッスル兄弟
2匹以上になるとだいたい変になる。
リアル猫バス?
身体の生成は苦手っぽい。

GoogleColabを使って試す(YouTube)

GoogleColabデモの説明はNNabla公式YouTubeで詳しく説明されております。是非ご覧ください。

コメント

タイトルとURLをコピーしました