PyTorch版StyleGAN2の編集

Introduction

StyleGANをはじめとする画像生成モデルが注目を集めている。生成モデルはノイズから画像への写像と捉えられるが、この逆写像を求める試みはGAN Inversionと呼ばれ、様々な手法が提案されている[1]。 GAN Inversionにおける実装では、前提とする生成モデルにPyTorch版のStyleGAN2がしばしば用いられる[2]。このモデルでは、一部の関数をCUDAのカーネルで置き換えて高速化しているが、CPUのみの環境での実行を目指すと、この部分の実行にエラーが付きまとう。
本稿では、とりあえず通常のPython & PyTorchでこの処理を実施して、エラーを抑制する方法について検討する。

Method

編集個所は2か所ある。大抵のリポジトリでは、modelsの下にstylegan2が置かれていて、stylegan2/opディレクトリ内に原因の箇所がある。

一つ目はfused_act.pyで、FusedLeakyReLUと呼ばれる処理が記述されている。 [3]によれば、この活性化関数の処理は自明なので、オリジナルの定義を変更する。
二つ目はupfirdn2d.pyで、こちらはファイル内に既にPython & PyTorch版の実装が用意されているので、こちらを利用する。
両方とも、ファイルの上の方のmodule_pathやfusedの定義はコメントアウトしておく。

<追記>色々と編集個所をこの記事で示していたが、[4]の記述に従ってうまく実行できたのでそちらを参照してほしい。特に、upfirdn2dのnativeはそのままではうまく動かないので、[4]に従う。

Experiments

[1]で紹介されているHFGI[5]と呼ばれるInversionモデルを利用してみる。このモデルでもPyTorch版のStyleGAN2が利用されているため、CPUでデモを動かすには上記の作業が必要である。

Dockerfile CPUでデモを動かすことを考えて、以下のようなDockerfileを準備した。

FROM pytorch/pytorch:1.6.0-cuda10.1-cudnn7-runtime
RUN pip install --upgrade pip
RUN pip install matplotlib

その他の準備

Methodのセクションの2オペレーションの編集
scripts/inference.pyの143行目をdevice = "cpu"に変更
[6]に従って学習済みモデルをダウンロードしてcheckpoint/ckpt.ptに配置する

以上の作業により、test_imgs/以下の12枚の画像が入力される。 HFGIではStyleの編集も簡単にできるため、以下で試みる。

inference.shを編集し、引数を変更する：--edit_attribute='smile'
inference.shを編集し、編集の強さを変更する。デフォルトでは0なので明示的に指定しないとスタイル編集されない：--edit_degree=1.0

以上が完了したのち、bash inference.shを実行することで、experiment/inference_results/ディレクトリに結果が出力される。

References

[1] GitHub - weihaox/awesome-gan-inversion: A collection of resources on GAN inversion.
[2] GitHub - rosinality/stylegan2-pytorch: Implementation of Analyzing and Improving the Image Quality of StyleGAN (StyleGAN 2) in PyTorch
[3] Definition of FusedLeakyRelu · Issue #77 · rosinality/stylegan2-pytorch · GitHub
[4] Bug in upfirdn2d_native and native pytorch version of it. · Issue #81 · rosinality/stylegan2-pytorch · GitHub
[5] GitHub - Tengfei-Wang/HFGI: CVPR 2022 HFGI: High-Fidelity GAN Inversion for Image Attribute Editing
[6] GitHub - Tengfei-Wang/HFGI: CVPR 2022 HFGI: High-Fidelity GAN Inversion for Image Attribute Editing