GCPで金をかけずにStylegan2のFinetuning #DeepLearning

Finetuning (transfer learning)は学習済みのモデルを元にして、希望のターゲットについて学習させ直す技術。正しいモデルを選ぶと１から学習させるよりも少ない学習データで圧倒的に早く望みの結果を得られ、学習にかかる時間とコストを削減できるぞ。

今回は人間の顔を生成するモデルをFinetuneして、anime-facesを生成させるようにする。
また、マシンも金もないので、GCPのVMを使ってなるべくコストを抑える。

参考

https://github.com/NVlabs/stylegan2
https://towardsdatascience.com/stylegan-v2-notes-on-training-and-latent-space-exploration-e51cf96584b3

環境

[GCP]
Machine type:
　custom (8 vCPUs, 50 GB memory)
GPUs:
　1 x NVIDIA Tesla K80
Disc:
　200GB
Framework：
　TensorFlow Enterprise 1.15

※学習途中でエラーになるので、メモリとディスクはケチっちゃダメ
※gpuのquoraを増やすのを忘れないように
※Frameworkに気をつけて、Stylegan2はtf2.0では動かないぞ
※お値段$0.667/hour

準備

1.元となるモデルをダウンロード

以下から faces (FFHQ config-e 256x256) をダウンロード
https://reposhub.com/python/deep-learning/justinpinkney-awesome-pretrained-stylegan2.html

gdown https://drive.google.com/uc?id=1BUL-RIzXC7Bpnz2cn230CbA4eT7_Etp0

2.データセットをダウンロード

今回はkaggleから以下を使用
https://www.kaggle.com/splcher/animefacedataset

pip install -q kaggle
cp kaggle.json ~/.kaggle/
chmod 600 ~/.kaggle/kaggle.json
kaggle datasets download -d splcher/animefacedataset
unzip animefacedataset.zip -d animefacedataset

3.Stylegan2のリポジトリをClone

git clone https://github.com/shawwn/stylegan2.git

4.データセットの画像を学習用に加工する

Stylegan2の学習用画像は正方形かつ辺のサイズが2の累乗でなくてはいけないため、事前に画像のサイズを加工しておく。今回はPythonのスクリプトを使い、256 * 256のサイズに加工した。

from PIL import Image
import PIL
import os
import glob
filenames = glob.glob("~/animefacedataset/images/*.jpg")
stylegan_dir = '~/stylegan2/animefacedataset/'
base_width = 256
for filename in filenames:
  new_filename = stylegan_dir + filename.split('/')[-1]
  image = Image.open(filename)
  width_percent = (base_width / float(image.size[0]))
  hsize = int((float(image.size[1]) * float(width_percent)))
  image = image.resize((base_width, hsize), PIL.Image.ANTIALIAS)
  if image.size[0] == image.size[1] == base_width:
    image.save(new_filename)

5.データセットを作成する

画像をmulti-resolution TFRecordsに変換する
リポジトリに含まれているdataset_tool.pyを使って簡単に変換できる。
詳細は以下
https://github.com/NVlabs/stylegan2#preparing-datasets

python dataset_tool.py create_from_images ~/stylegan2/datasets/animedataset ~/animefacedataset

[エラーあるある]

以下は正方形でない画像が含まれていた場合に起きるエラー

Loading images from "./animefacedataset"
Creating dataset "./datasets/animedataset"
Added 16 images.
Traceback (most recent call last):
  File "dataset_tool.py", line 642, in <module>
    execute_cmdline(sys.argv)
  File "dataset_tool.py", line 637, in execute_cmdline
    func(**vars(args))
  File "dataset_tool.py", line 526, in create_from_images
    tfr.add_image(img)
  File "dataset_tool.py", line 78, in add_image
    assert img.shape == self.shape
AssertionError

6.Finetuning用にスクリプトを改造

元のrun_training.pyを改造してFinetuningを行えるようにする。

run()に以下を書き足す

run_training.py

    train.resume_pkl = "/home/kitigai/ffhq-256-config-e-003810.pkl"  # 1でダウンロードした学習元モデル, None = train from scratch.
    train.resume_kimg = 3810.0  # 学習元モデルの学習進捗 別に0.0でも良いと思う.
    train.resume_time = 0.0  # 学習元モデルの学習時間 分からないし別に0.0でも良いと思う.

元のmetrics/frechet_inception_distance.pyがやたらでかい.pklをダウンロードしようとして"Google Drive quota has been exceeded"になるのでローカルにダウンロードしてきた.pklを使用するように変更

frechet_inception_distance.py

    def _evaluate(self, Gs, Gs_kwargs, num_gpus):
        minibatch_size = num_gpus * self.minibatch_per_gpu
        #inception = misc.load_pkl('https://nvlabs-fi-cdn.nvidia.com/stylegan/networks/metrics/inception_v3_features.pkl')
        inception = misc.load_pkl('/home/kitigai/inception_v3_features.pkl', 'inception_v3_features.pkl')
        activations = np.empty([self.num_images, inception.output_shape[1]], dtype=np.float32)

7.学習開始

コマンドの詳細は以下
https://github.com/NVlabs/stylegan2#training-networks

python run_training.py --num-gpus=1 --data-dir=datasets --config=config-e \
  --dataset=animedataset --mirror-augment=true

[エラーあるある]

--data-dirや--datasetに渡すパスに"/"や"."が入っているとエラーになる

# ダメなやつ
--data-dir=./datasets 
# ダメなやつ２
--dataset=/home/kitigai/stylegan2/datasets/animedatasets

半端なメモリだとエラーになる

RAM:30GB のVMを使っていたら以下のエラーになったので50GBまで増やした。ケチってはいけない

Traceback (most recent call last):
  File "run_training.py", line 192, in <module>
    main()
  File "run_training.py", line 187, in main
    run(**vars(args))
  File "run_training.py", line 120, in run
    dnnlib.submit_run(**kwargs)
  File "/home/myuser/stylegan2-master/dnnlib/submission/submit.py", line 343, in submit_run
    return farm.submit(submit_config, host_run_dir)
  File "/home/myuser/stylegan2-master/dnnlib/submission/internal/local.py", line 22, in submit
    return run_wrapper(submit_config)
  File "/home/myuser/stylegan2-master/dnnlib/submission/submit.py", line 280, in run_wrapper
    run_func_obj(**submit_config.run_func_kwargs)
  File "/home/myuser/stylegan2-master/training/training_loop.py", line 341, in training_loop
    metrics.run(pkl, run_dir=dnnlib.make_run_dir_path(), data_dir=dnnlib.convert_path(data_dir), num_gpus=num_gpus, tf_config=tf_config)
  File "/home/myuser/stylegan2-master/metrics/metric_base.py", line 151, in run
    metric.run(*args, **kwargs)
  File "/home/myuser/stylegan2-master/metrics/metric_base.py", line 67, in run
    self._evaluate(Gs, Gs_kwargs=Gs_kwargs, num_gpus=num_gpus)
  File "/home/myuser/stylegan2-master/metrics/frechet_inception_distance.py", line 54, in _evaluate
    labels = self._get_random_labels_tf(self.minibatch_per_gpu)
  File "/home/myuser/stylegan2-master/metrics/metric_base.py", line 140, in _get_random_labels_tf
    return self._get_dataset_obj().get_random_labels_tf(minibatch_size)
  File "/home/myuser/stylegan2-master/metrics/metric_base.py", line 121, in _get_dataset_obj
    self._dataset_obj = dataset.load_dataset(data_dir=self._data_d![3810.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/896087/955d309d-eff5-535a-ba30-c0b6a31b50b3.png)
ir, **self._dataset_args)
  File "/home/myuser/stylegan2-master/training/dataset.py", line 192, in load_dataset
    dataset = dnnlib.util.get_obj_by_name(class_name)(**kwargs)
  File "/home/myuser/stylegan2-master/training/dataset.py", line 86, in __init__
    self._np_labels = np.zeros([1<<30, 0], dtype=np.float32)
MemoryError: Unable to allocate 0 bytes for an array with shape (1073741824, 0) and data type float32