StyleGAN2 ADAで作る見たことないF1車体画像

はじめに

本記事はABEJAアドベントカレンダー2021の14日目の記事です。

私はF1レースのファンで、去年のアドベントカレンダーでもF1の面白さを予測するというテーマで記事を投稿しました(いやー、今シーズンは大接戦で最終戦の最後の周回まで手に汗握る展開でしたね。ありがとう、ホンダF1)。

今年はThis Person Does Not Existなどでも有名なGANベースの生成モデルであるStyleGAN2を使って、この世界にまだ存在しないF1カーの車体画像を作ってみたいと思います。特に今回は、StyleGAN2からデータ拡張(Data Augmentation)部分が改良された、数千枚程度の学習データ画像でもうまくいくというStyleGAN2 ADAを試してみました。

先に結果をお見せしちゃうと、学習後以下のような品質の画像の生成をすることができました。この例では、左から2004年のルノー→レッドブル→フェラーリのような感じで潜在空間と画像のマッピングを学習していることが見て取れます。

データ準備

StyleGAN2 ADAで学習データとして使える画像には、以下のような制限があります。

• 正方形であること
• 一辺の長さが$2^n$(例：32, 64, 128, 256, ...)ピクセルになっていること

また学習のしやすさを考えると、なるべく一台の車体が画像の中央に配置されているのが理想的です。しかし、残念ながらそのような都合の良いF1車体画像データセットは、私が探した範囲では見つかりませんでした。

F1画像を掲載しているニュースサイトや壁画像サイトは結構あるのですが、何千枚もの単位で上記の条件を満たすようなデータは当然ありません。そこで、Sketchfabという3Dモデルの共有・売買サイトで、TheoDevF1氏が公開しているF1車体の3Dモデルを使用させてもらうことを考えました。

以下のような3DモデルをグリグリとSketchfabサイト上で動かしながらを動画キャプチャし、その動画から256x256サイズの画像を6000枚程度作成しました（以下は同氏が無償で公開している2004年シーズンの車体モデルの一部です）。

StyleGAN2 ADAによる学習

実行環境の準備

• StyleGAN2は、Tensorflow 1.xの環境が必要です。私はUbuntu上に、こちらのインストラクションを参考にTensorflow 1.15の環境を作ったのですが、学習を実行させると以下のような警告メッセージが出てきたので1.14の方が良いようです。

WARNING: Using slow fallback implementation for inter-GPU communication.
Please use TensorFlow 1.14 on Linux for optimal training performance.

• NVIDIA社による以下のStyleGAN2 ADA実装（StyleGAN2 with adaptive discriminator augmentation - Official Tensorflow implemtation）をgit cloneするなどして、実行環境に配置します。
  • https://github.com/NVlabs/stylegan2-ada

データ変換

• Tensorflowでの処理のためには入力画像をTFRecord形式に変換してやる必要があり、StyleGAN2 ADAによって提供される以下のユーティリティコマンドを実行します。

$ python dataset_tool.py create_from_images \
    datasets/formula1/tfrecords \  # 出力TFRecordディレクトリ
    datasets/formula1/raw_images   # 入力画像ディレクトリ

学習の実行

• 以下のコマンドにより、↑で準備したTFRecordファイルを入力として学習を実行します。
  • ssh経由でリモートマシン上で学習を行う場合、ネットワーク切断などの外乱によりコマンド実行が中断されないようnohupを指定してやるのが良いと思います。

$ nohup python train.py \
  --data=datasets/formula1/tfrecords \ # 入力TFRecordディレクトリ
  --outdir=training-runs \             # ログ出力ディレクトリ
  --gpus=2 \                           # 使用するGPUの数
  --kimg=500                           # 学習に用いる合計画像枚数[千枚単位]

• 学習中は、定期的にモデルのスナップショットがpklファイルに保存されます。学習を再開したい場合には、以下のように保存されたpklファイルをresumeオプションで指定するだけでOKです。
  • ちなみにこのresumeオプションは、StyleGAN2で訓練された任意の学習済みモデル（例えば、awesome-pretrained-stylegan2など）を指定してモデルのファインチューンするのにも使用できます。

$ nohup python train.py \
  --data=datasets/formula1/tfrecords \ 
  --outdir=training-runs \             
  --gpus=2 \                           
  --kimg=500 \
  --resume=training-runs/00000-tfrecord-auto2-kimg500/network-snapshot-000500.pkl  # 学習再開したいモデルを指定

実験結果

学習の様子

今となっては年代物の、2015年製GeForce GTX 1080 GPU * 2枚で学習をひたすら行います！

GANにおいて、生成画像と元画像の類似度を図る指標としてよく用いられるFrechet Inception Distance (FID)が、累計の学習画像枚数に応じてどのように推移するかをプロットしたのが以下の図になります。この数値が低ければ低いほど、生成画像が元画像に似ているということになります。なお、累計230万枚ほどを学習させるのに上記のGPUで丸2日ほどかかりました。

150万枚くらい学習したところでほぼFIDは収束して、最終的にベストのFID(=34.6419)は200万枚分学習したところで得られました。

以下に学習枚数毎のFIDと生成画像のサンプルを示します。個人的には学習枚数が100万枚に到達した頃に、タイヤが丸くなったことが非常に嬉しかったです笑。

学習画像枚数 [100万枚]	0.2	0.327	0.5
FID	345.5018	167.2883	91.3105
生成画像

| 学習画像枚数 [100万枚] | 1.0 | 1.532 | 2.0 |
|:-----------:|:-----------:|:------------:|:------------:|:------------:|:------------:|:------------:|
| FID | 52.4651 | 37.9323 | 34.6419 |
| 生成画像 | | | |

生成画像

以下は、200万枚学習時に生成された画像の一部です。タイヤがひしゃげていたり、ディテールが潰れてしまっているものも一部ありますが、ほとんどの画像でF1車体っぽい画像が生成できているように感じます。ただ、期待したような「これは見たこと無い！」というような驚きが感じられるような車体のカラーリングや形を生成するまでには至りませんでした。

冒頭でも示した潜在空間と画像のマッピングアニメーション生成には、こちらのBlog記事で紹介されていたコードを使わせてもらいました。F1車体のような複雑な造形の物体でも連続的に潜在空間と画像のマッピングが学習されていることが分かります。

まとめ

StyleGAN2 ADAを使ってF1カーの車体画像の生成を試みました。正直に言うと、表題の「見たことの無いF1車体画像を作る」レベルまでには至りませんでしたが、F1車体のような複雑な形でも破綻なく生成できるモデルを数千枚の画像から学習できる事自体には驚きました。

改めてこのような優れた技術を開発し、共有してくれた先人達に感謝したいです！

お知らせ

現在ABEJAでは一緒にAIの社会実装を進める仲間を募集しています。
ABEJA Advent Calendar 2021を読んで少しでもいいね！と思ったら、まずはお話を聞きに来てください。
【現在募集中の職種】 はこちらから確認できます。ご応募をお待ちしております。