自前PCにStyleGAN2の学習環境を作る(Win10+GTX1080ti)

StyleGAN2の学習環境を自前PC(Windows10+GTX1080ti)で構築した時の手順を紹介します。
ゲーミングPCなどローカル環境で学習させたい方の参考になればと思います。
GPUのメモリが16GBに満たない場合、1024x1024の解像度では出力できません。

ちなみにこちらの勉強会で発表した内容です。

全体の流れ

• ソフトのインストール
• ソース修正
• 動作確認
• 学習

ソフトのインストール

Visual Studio

Visual Studio 2017 Communityをインストールします。
※Visual Studio 2019だけが入った環境では動作しませんでした。適切なライブラリを入れればOKかもしれません。

インストール後、環境変数PATHに"vcvars64.bat"と"cl.exe"のフォルダまでのパスを通しておきます。

例えば下記のようになります。

• "C:\Program Files (x86)\Microsoft Visual Studio\2017\Community\VC\Auxiliary\Build\vcvars64.bat"
• "C:\Program Files (x86)\Microsoft Visual Studio\2017\Community\VC\Tools\MSVC\14.16.27023\bin\Hostx64\x64"

CUDA、cuDNN

CUDA 10.0とcuDNNをインストールします。
※CUDA 10.1では動作しませんでした。

インストール後の確認として、コマンドプロンプトで下記を実行します。

nvcc test_nvcc.cu -o test_nvcc –run

文字コードに関する警告が出ますが、最終的に下記が表示されればOKです。

CPU says hello.
GPU says hello.

Pythonライブラリ

Anacondaで仮想環境を構築し、下記のライブラリをインストールしました。

conda install numpy=1.14.3
conda install tensorflow-gpu=1.14
conda install pillow
conda install requests

私の実行時にはpillowは7.0.0、requestsは2.22.0が入りました

インストール後の確認として、ライブラリをインストールした仮想環境のコマンドプロンプトで下記を実行します。
エラーが出なければOKです。

python -c "import tensorflow as tf; tf.enable_eager_execution(); print(tf.reduce_sum(tf.random_normal([1000, 1000])))“

ソース修正

コンパイラのパスを環境に合わせる

custom_ops.pyの27行目あたりにコンパイラのパスが指定されているので、これを自分の環境に合わせます。
例）

stylegan2\dnnlib\tflib\custom_ops.py

compiler_bindir_search_path = [
    ～
    'C:/Program Files (x86)/Microsoft Visual Studio/2017/Community/VC/Tools/MSVC/14.16.27023/bin/Hostx64/x64',
    ～]

学習のさせ方に合わせて設定（任意）

出力する解像度を指定する場合、run_training.py 61行目のdataset_argsにresolutionを追加指定します。
2のべき乗の値を指定する必要があるのと、学習用画像の解像度よりも小さい値にする必要があります。

stylegan2\run_training.py

dataset_args = EasyDict(tfrecord_dir=dataset, resolution=64)

学習経過のモデルと出力画像のスナップショットを取る頻度を指定する場合、run_training.py 51行目の値を変更します。

stylegan2\run_training.py

train.image_snapshot_ticks = train.network_snapshot_ticks = 1

動作確認

以上までで動作するはずです。コマンドプロンプトで下記を実行します。
既存モデルで画像が出力されればOKです。

python run_generator.py generate-images --network=gdrive:networks/stylegan2-ffhq-config-f.pkl --seeds=66,230,389,1518 --truncation-psi=1.0

学習

あとは学習で使う画像を集めたり、画像を調整したりし、学習を実行します。

学習時には画像をtfrecord形式にする必要があります。コマンドプロンプトで下記を実行します。

python dataset_tool.py create_from_images tfrecord形式出力先フォルダ 入力元画像フォルダ

学習実行時にはconfigセット(a～f)を指定できます。
aとf(デフォルト)しか試していませんが、同じ学習時間でもaよりもfの方が本物っぽい画像が出力されました。

参考

StyleGAN2 — Official TensorFlow Implementation