styleGANで独自データを学習する

2021.05.16
Python

以前、アニメキャラをDCGANに学習させてみたことがあるのですが、ふと人間以外のキャラクタだとどうなるのかと思ったので、今度はstyleGANでポケモンの画像を学習させてみました。

使用するコードは下記で公開されているものです。学習済みのモデルも公開されており、非常にハイクオリティな人物や猫の画像がすぐに生成できます。

GitHub

NVlabs/stylegan

https://github.com/NVlabs/stylegan

StyleGAN - Official TensorFlow Implementation. Contribute to NVlabs/stylegan development by creating an account on GitHub.

1. 学習画像準備
2. コード修正
3. 学習
4. 学習結果

学習画像準備

今回は、まず学習用の画像を準備します。なお、ポケモンの画像はデータセットとしてネットにあったものを使用しました。著作権的に問題ありそうなのでURLは載せません。なお、注意点として解像度は縦横が同じ、かつ2の累乗である必要があるので、事前にサイズを変更しておきます。

画像はdataset_tool.pyを使って変換する必要があります。最初に、フォルダを作成してそこに学習に使いたい画像を保存します。次にdataset_tool.pyを起動します。下記の例だとcutom-imagesにある画像を変換して、dataset/custom-datasetフォルダに保存することになります。

python dataset_tool.py create_from_images datasets/custom-dataset custom-images

1	python dataset_tool.py create_from_images datasets/custom-dataset custom-images

コード修正

画像を準備しただけでは、学習できないのでtrain.pyを修正します。

まず準備した画像を使用できるにようにフォルダの指定をします。なお、出力する画像の解像度も指定できます。下記のコードの場合、train.pyの起動時に引数として-customを指定することでcustom-datasetフォルダのデータを学習し、解像度128×128の画像を出力します。

desc += '-custom'; dataset=EasyDict(tfrecord_dir='custom-dataset', resolution=128)

1	desc += '-custom'; dataset=EasyDict(tfrecord_dir='custom-dataset', resolution=128)

また、同じファイルにあるtrain.total_kimgが学習回数のようなので、任意で変更します。

学習

学習は下記のようにtrain.pyを起動します。私のPCにGPUは1個しか無いので引数で1gpuと指定しています。

python train.py -custom -1gpu

1	python train.py -custom -1gpu

学習結果

4000回ほど学習させた結果、下記のような感じになりました。はっきりとした感じではなくモザイク画像のようになっています。これで学習時間は9時間くらいです。本当は数日かけてやるべきだと思いますが、とりあえず、ちゃんと学習できていることは分かりました。

人物は個人差はあるものの基本的な体や顔の構造は同じです。なので、特徴を抽出すればちゃんと人間としての画像が生成されます。ただ、ポケモンのように姿形が全く違うと、共通点も殆ど無いので、特徴量と呼べるような物が無いのではという気がします。やはりこれまでに無い全く新しい何かを作るのはディープラーニングでは難しそうです。