2025.01試行
ubuntu2204 RTX4090

styleGANを試してみよう

stylegan1,stylegan2のオリジナルの実装はtensorflow1系でdockerを使ってもうまく動かないなど難しい。

stylegan2-adaという、学習が安定するバージョンもtensorflow1系だが、pytorchの実装も存在する。

stylegan2-ada-pytorchなら動きそうだ

今回は、stylegan2-ada-pytorch　を試した。

docker環境とネイティブ環境での参考コードがありますが、環境構築の手間が無いdocker環境を使用します。

git clone https://github.com/NVlabs/stylegan2-ada-pytorch.git

cd stylegan2-ada-pytorch

サンプルコード動くか？

とりあえず、画像生成してみよう

docker build --tag sg2ada:latest .
./docker_run.sh python3 generate.py --outdir=out --trunc=1 --seeds=85,265,297,849 \
    --network=https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/metfaces.pkl

docker build --tag sg2ada:latest .でフォルダ内のDockerfileからイメージを作成します。
次のコマンドで、コンテナ内のpythonプログラムを実行して、4種類のシードから、画像を生成して、outフォルダに格納します。学習済みのpklを（初回）ダウンロードします。シードを変えてみて異なる画像が生成されることを確認してみる。

さて、dockerがどこのフォルダで実行されているのか確認してみると、

./docker_run.sh pwd
$ /scratch
./docker_run.sh ls
$ Dockerfile	 dataset_tool.py  generate.py  out .....など

内部では、/scratchフォルダがホームであることが分かります。（docker_run.shにあるとおりですが、）

潜在空間で元画像(生成画像)を探してみよう

ランダムシードで作成した画像だけど、こんな画像はどんな潜在空間にあるのだろう？

ここのサンプルコードで面白いのは、潜在空間への画像の投影(google翻訳）のところで、

コンテナ内で先ほど作成された画像out/seed0085.pngの潜在ベクトルを見つける過程は、

 ./docker_run.sh python3 projector.py --outdir=out2 --target=/scratch/out/seed0085.png --network=https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/metfaces.pkl

潜在空間を探索しながら元画像に近づいていく様子が分かります。(1000試行してロスが減少していく様子と外時の画像がmp4ファイルで保存される)

学習データと別のモデル(pkl)を使ってみるのも面白い。

 ./docker_run.sh python3 projector.py --outdir=out2 --target=/scratch/out/seed0085.png --network=https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/ffhq.pkl

まあ、ここまでは上手く行った。今日はここまで。