*Interp移植録 - 画像生成/ StyleGAN2-ADA (TflInterp)

0.Prologue

暇つぶしに、興味を引いた DNNアプリを *Interpに移植して遊んでいる。
本稿はその雑記&記録。

とある事情から、Tensorflowで実装された面白そうなモデルがないかと探していたところ、StyleGAN2 ADAと言う画像生成モデルが目に止まった。まぁ、画像生成といえば Stable-Diffusionがデファクトスタンダードな昨今、ちょい古な StyleGAN2は今更と言う感があるが、移植ボリュームが小さそうなのでやってみることにしたのだが……

1.Original Work

• "A Style-Based Generator Architecture for Generative Adversarial Networks"
  https://arxiv.org/abs/1812.04948

• Analyzing and Improving the Image Quality of StyleGAN
  https://arxiv.org/abs/1912.04958

NVIDIAの開発による StyleGAN2 ADAは、モデルへの入力(latants)を操作することで、生成される画像内の特徴パーツ - 例えば人の顔の画像の生成であれば、髪型や目、口のなど - のスタイルを個別にコントロールできる画像生成モデルだ。

モデルの構造は下図のようにMappingとSynthesisの二つのネットワークから成っており、その接合部の潜在変数空間が線形となるように訓練されているようだ。この潜在変数空間の線形性が特徴パーツの個別コントロールを可能にしているのだろう。

StyleGANは、初出の StyleGAN(Mar 2019)から改良が繰り返されており、ごく最近では "Text-to-image"タスクへの応用が試みられている - StyleGAN-T(Jan 2023)。

StyleGAN(初出) → StyleGAN2(ノイズ低減)
　→ StyleGAN ADA(データ拡張訓練)
　→ StyleGAN3(生成部改良) → -XL,-T

今回の移植に取り上げたモデルは StyleGAN2-ADAで、少ない訓練データでも精度よく学習できるようにデータ拡張(Data Augmentation)をトレーニングに取り入れたモデルだ。まぁ、移植する部分はその改良部分には触れない Generater(画像生成)部分だけなので、一つ前の StyleGAN2と何ら変わるところは無い…

StyleGAN2:

2.StyleGAN2-ADAの tfliteモデルの調達

冒頭で触れた通り、StyleGAN2-ADAはTensorflowで実装されている。故に、移植先のTflInterpに喰わす tfliteモデルを調達するのは楽勝と踏んでいたのだが……TF1だとぉぉ おまけに pretraindモデルは pickle形式?

結局、モデルのコードをTF1からTF2にコンバートし、さらに pickle形式の pretrainedモデルを Savedmodelを経由して tfliteに変換するスクリプト群を用意することした。その詳細は割愛するが、下の手順に従えばpickleから tfliteモデルを調達できる。この手順では、StyleGAN2-ADAのMappingとSynthesisの二つのネットワークそれぞれの tfliteが得られる - afhqdog.mapping.tflite, afhqdog.synthesis.tflite.

# Tensorflow v2にコンバートしたStyleGAN2-ADAをクローンする
git clone https://github.com/shoz-f/stylegan2-ada.git
cd stylegan2-ada
# pickleから Savedmodelに変換する
python pkl2savedmodel.py https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada/pretrained/afhqdog.pkl afhqdog
# Savedmodelから tflite に変換する
python export_tflite.py -s mapping afhqdog
python export_tflite.py -s synthesis afhqdog 

※ちょこっと試してみたところ、StyleGAN2の pickle形式モデルも上の手順で tfliteに変換できるようだ

3.TflInterp用のLivebookノート

では、拙作のTflInterpを利用して StyleGAN2-ADAを遊んでみよう。MappingとSynthesisそれぞれを実行するモジュールを用意し、それらを組み合わせて使うように設計した。

まず最初に、Mix.installの依存リストの記述は次の通り。

File.cd!(__DIR__)
# for windows JP
System.shell("chcp 65001")
System.put_env("NNCOMPILED", "YES")

Mix.install([
  {:tfl_interp, "~> 0.1.13"},
  {:cimg, "~> 0.1.19"},
  {:kino, "~> 0.7.0"},
  {:nx, "~> 0.5.3"}
])

Mappingモジュールには、乱数で入力潜在変数(latants)を生成し内部潜在変数を返す dlatants/0と、f32:{512}の tensorを入力潜在変数(latants)として受け取り内部潜在変数を返す dlatants/1を用意した。

[モデル・カード]

• inputs:
  [0] f32:{1,512} - 0.0～1.0の範囲の値をとる512個の要素が並んだ潜在変数(ベクトル)
• outputs:
  [0] f32:{1,16,512} - 内部潜在変数

defmodule StyleGAN2ADA.Mapping do
  alias TflInterp, as: NNInterp

  use NNInterp,
    model: "./model/afhqdog.mapping.tflite",
    url: "https://github.com/shoz-f/tfl_interp/releases/download/0.0.1/afhqdog.mapping.tflite",
    inputs: [f32: {1, 512}],
    outputs: [f32: {1, 16.512}]

  def dlatants() do
    for _ <- 1..512, into: "" do <<:rand.normal()::little-float-32>> end
    |> Nx.from_binary(:f32)
    |> dlatants()
  end

  def dlatants(latants) do
    # preprocess
    input0 = Nx.to_binary(latants)

    # prediction
    output0 =
      session()
      |> NNInterp.set_input_tensor(0, input0)
      |> NNInterp.invoke()
      |> NNInterp.get_output_tensor(0)

    # postprocess
    Nx.from_binary(output0, :f32) |> Nx.reshape({16, 512})
  end
end

Synthesisモジュールは、Mappingモジュールにより生成された内部潜在変数を受け取り、CImg {512,512,1,3}形式の画像を返す(image/1)

[モデル・カード]

• inputs:
  [0] f32:{1,16,512} - 内部潜在変数
• outputs:
  [0] f32:{1,3,512,512} - 画素値に -1.0～1.0の値をとるRGB画像tensor

defmodule StyleGAN2ADA.Synthesis do
  @width 512
  @height 512

  alias TflInterp, as: NNInterp

  use NNInterp,
    model: "./model/afhqdog.synthesis.tflite",
    url: "https://github.com/shoz-f/tfl_interp/releases/download/0.0.1/afhqdog.synthesis.tflite",
    inputs: [f32: {1, 16, 512}],
    outputs: [f32: {1, 3, @height, @width}]

  def image(dlatants) do
    # preprocess
    input0 = Nx.to_binary(dlatants)

    # prediction
    output0 =
      session()
      |> NNInterp.set_input_tensor(0, input0)
      |> NNInterp.invoke()
      |> NNInterp.get_output_tensor(0)

    # postprocess
    CImg.from_binary(output0, @width, @height, 1, 3, [{:range, {-1.0, 1.0}}, :nchw])
  end
end

4.デモンストレーション

StyleGAN2ADA.Mapping, StyleGAN2ADA.Synthesisを起動する。

# TflInterp.stop(StyleGAN2ADA.Mapping)
StyleGAN2ADA.Mapping.start_link([])

# TflInterp.stop(StyleGAN2ADA.Synthesis)
StyleGAN2ADA.Synthesis.start_link([])

Mapping.dlatants/0で２つのランダムな内部潜在変数を生成し、その間を5等分する点を求め、それぞれの点に対応する画像を生成してみよう。きっと…

a = StyleGAN2ADA.Mapping.dlatants()
b = StyleGAN2ADA.Mapping.dlatants()
# 5等分の間隔
delta = Nx.subtract(b, a) |> Nx.divide(5)

Enum.map(0..5, fn i ->
  Nx.add(a, Nx.multiply(delta, i))
  |> StyleGAN2ADA.Synthesis.image()
  |> CImg.display_kino(:jpeg)
end)
|> Kino.Layout.grid(columns: 6)

何度かリトライを要したが、a～bに順に Morphする画像列が生成できた こんなことができるのは、内部潜在変数空間が線形なっているお陰だね。

5.Epilogue

Diffusion系の画像生成の台頭により、GANはすっかりと影を潜めてしまった感がある。なんでも、GANはトレーニングが難しいからだそうだ。とは言え、画像の生成速度は結構速いので、何か応用できるフィールドがないものかと思う。

それにしても、モデルを MappingとSynthesisの２つに分割し、その結合部の内部潜在変数空間を線形に導くというアイデアは、どうすりゃ思いつくのだろう素晴らしい

Appendix

• TflInterpのノート
  https://github.com/shoz-f/tfl_interp/blob/main/demo_StyleGAN2-ADA/StyleGAN2ADA.livemd

• StyleGAN-ADAの tfliteモデル調達手順
  https://github.com/shoz-f/tfl_interp/blob/main/demo_StyleGAN2-ADA/Export_StyleGAN2_ADA(afhqdog).ipynb