Linux上でStarGANの実装をしてみた

今回は、StarGANの実装をしてみました。
基本的に、githubで公開されているコードをもとに実装を行っていきます。
本ページでは、軽い論文の説明と実装をしていきます。
自分のデータセットを用いて適用する回は、次回やっていきたいと思います。

StarGANについて超簡単に
Linux上での実装

簡単ではありますが、上の２項目に従って説明します。

StarGANについて超簡単に

論文：https://arxiv.org/abs/1711.09020

StarGANとは、異なるドメイン間への変換を可能とするGenerative Adversarial Netwok (GAN)の一種であるCycleGANと、GANに多クラス分類の学習を加えたAC-GANを組み合わせたものです。

CycleGANは、2つのドメイン間でのドメイン変換しか行えません。
2つ以上のk個のドメイン間でのドメイン変換を実現するためには、k(k-1)個のGeneratorを学習しなければなりません。
これは、事実上可能ではありますが、面倒くさいです。

この画像は、論文内のものですが、k=4だとすると、4*3=12個のGeneratorが必要だということを表しています。

そんな問題を解決するために、StarGANでは、1つのGeneratorで複数のドメインの変換を行うアルゴリズムを導入しています。

5つのドメインがある場合、上の画像のようにあらわされます。
星形をしていますね。だからStarGANなんでしょう。なんか厨二心をくすぐられます。

丁寧な説明が他の記事ありましたので、共有いたします。（これ以上の解説は自分のはできない）

Linux上での実装

公開コード
https://github.com/yunjey/StarGAN

実装環境

Ubuntu 18.04 LTS
Python 3.6
PyTorch 0.4.0
Tensorflow 1.4.0

公開データセットで実装

まずは、任意のディレクトリにgitをクローンします。
続いて、StarGAN/のディレクトリに移動します。
今回は、論文でも使用されていたCelebA datasetとRaFDいうデータセットをダウンロードしていきます。

$ git clone https://github.com/yunjey/StarGAN.git
$ cd StarGAN/
$ bash download.sh celeba

ダウンロード完了です。

  .
  .
  .
  inflating: ./data/celeba/images/072137.jpg  
  inflating: ./data/celeba/images/027742.jpg  
  inflating: ./data/celeba/images/188764.jpg  
  inflating: ./data/celeba/list_attr_celeba.txt

こちらのCelebA datasetというデータセットは、こちらで公開されております。

このデータセットには、有名人のカラー顔画像が178×218ピクセルで202,599枚含まれています。
これに加えて、それぞれの画像に対して、40種類のattribute (属性) が付与されています。
例えば、Black_Hair、Blond_Hair、Brown_Hair、Male、Youngなどなどです。
もっと知りたい人はこちらを参照ください。

ダウンロードディレクトリには、上のような画像とlist_attr_celeba.txtという各画像に対する属性が記されているtxtファイルが保存されます。
こちらの中身は、

[画像フォルダ名] 1 1 -1 1 -1 1 1 ... 1 -1

となっています。1 or -1は属性の数の分（40個）あります。
1のとき、その属性であり、-1のとき、その属性ではないことを表します。
こちらは、自分のデータセットで学習させる際に生成する必要があるので覚えておいてください。

学習

まずは、CelebA datasetを使ってStarGANを学習していきます。

$ python main.py --mode train --dataset CelebA --image_size 128 --c_dim 5 \
               --sample_dir stargan_celeba/samples --log_dir stargan_celeba/logs \
               --model_save_dir stargan_celeba/models --result_dir stargan_celeba/results \
               --selected_attrs Black_Hair Blond_Hair Brown_Hair Male Young

コマンドについて説明します。
--mode : モード
--dataset : 用いるデータセット
--image_size : 画像サイズ
--c_dim : attribute (属性) のクラスの数
--sample_dir : サンプルが保存してあるフォルダ
--log_dir : ログを保存するフォルダ
--model_save_dir : モデルを保存するフォルダ
--result_dir : 結果を保存するフォルダ
--selected_attrs : 学習したいattribute (属性)
という感じです。

テスト

続いて、テストをしていきます。

$ python main.py --mode test --dataset CelebA --image_size 128 --c_dim 5 \
               --sample_dir stargan_celeba/samples --log_dir stargan_celeba/logs \
               --model_save_dir stargan_celeba/models --result_dir stargan_celeba/results \
               --selected_attrs Black_Hair Blond_Hair Brown_Hair Male Young

簡単にできてしまいます。

続いて、RaFDのデータもダウンロードしたい、、、と思ったのですが、申請が必要でした。
おそらく、申請できても、こちらのページにも掲載することはできないと思います。

学習するのが面倒くさいという方のためにも、学習済みネットワークをダウンロードしていきます。

$ bash download.sh pretrained-celeba-128x128

ダウンロードされたモデルは、./stargan_celeba_128/modelsに保存されます。
つづいて、以下のコマンドで学習済みネットワークを使って画像変換をしていきます。

$ python main.py --mode test --dataset CelebA --image_size 128 --c_dim 5 \
                 --selected_attrs Black_Hair Blond_Hair Brown_Hair Male Young \
                 --model_save_dir='stargan_celeba_128/models' \
                 --result_dir='stargan_celeba_128/results'

という感じに超簡単です。

./stargan_celeba_128/resultsに変換画像が保存されます。
これで、公開データセットでの簡単な実装は終了。
次回は、自分で用意したデータセットに適用していきます。

参考資料

論文：https://arxiv.org/pdf/1703.10593.pdf
Github：https://github.com/yunjey/StarGAN
CelebA：http://mmlab.ie.cuhk.edu.hk/projects/CelebA.html
7変化! AIで顔を変えよう：https://qiita.com/uyuutosa/items/5dd9b2b80a9aaca67050