少ない画像データでGANを学習する ~Data-Efficient GANs with DiffAutment~

機械学習 画像処理

以下のレポジトリに公開されているコードを用いて数百枚程度の少ない画像でGANを学習できたので、二番煎じ感がありますが備忘録がてら手順を記事にしておきます。

github.com

環境

学習手順

以下では上記レポジトリをcloneしてある状態を仮定しています。

Dockerイメージの作成：

cd data-efficient-gans/DiffAugment-stylegan2-pytorch
sudo docker build -t stylegan:0.0 .

訓練データの作成：

解像度が512の画像データセットの場合です。

※いろいろオプションがあるようなので、helpも確認してみるといいかもしれません。

# data-efficient-gans/DiffAugment-stylegan2-pytorch以下で
python dataset_tool.py --source <訓練画像が含まれるディレクトリ> --dest dataset_name.zip --width 512 --height 512

訓練：

学習にまとまった時間がかかることを考えてバックグラウンドで学習を回したかったため、以下のようなシェルスクリプトを書きました。(以下ではrun_background.shという名前で作成したと仮定。)

#!/bin/bash

cd /stylegan
python -u train.py \
    --outdir=output_directory_name --data=dataset_name.zip --gpus 1 \
    --kimg 25000 --mirror True --gamma 1 --batch 64 1>out.out 2>&1

続いてdockerコンテナをバックグラウンドで起動して実行します。

sudo docker run -d --gpus all -v <DiffAugment-stylegan2-pytorchの絶対パス>:/stylegan \
    --shm-size=16g stylegan:0.0 bash -c "/stylegan/run_background.sh"

学習コードの標準出力は data-efficient-gans/DiffAugment-stylegan2-pytorch以下のout.outに、学習済みの重みファイルなどはoutput_directory_name以下に出力されます。

gradient accumulationについて：

上の例ではバッチサイズを64にしていますが、解像度512の画像ではGPUメモリが足りずに1ステップでは処理できません。そこでgradient accumulationを利用したいのですが、デフォルトの学習コードではaccumulation stepを外部から明示的に指定できるようにはなっていませんでした。

そこで訓練コードの一部を以下のように修正することで対応します。

# 旧：args.batch_gpu = batch // gpus を以下のように変更
args.batch_gpu = 8 # こちらの数値はお好みで

ハイパーパラメータについて

主に変更することになるハイパーパラメータは、学習の長さ（利用するのべ画像枚数）であるkimgと、正則化係数gamma、バッチサイズbatch（以下のデフォルト値の辞書ではmbが相当）あたりになるかと思います。

king, gammaについてはデータセットや画像の解像度によって望ましい値が異なるようですが、学習コードのデフォルト値が参考になるかもしれません。

※ちなみに上で紹介した学習実行例でのパラメータ設定は、私が1000枚程度で解像度が512の画像データセットで学習を行った際のものです。