NVIDIAの訓練済みモデルで遊ぶ

NVIDIA は訓練済みモデルを公開している

NVIDIA は最新のモデルの実装や、訓練済みモデルを公開しています。
カタログページ では、DeepLearning に限らず docker のコンテナなど分野問わず研究成果を公開していて見ていて楽しいです。
また、github の NVlabs では、実装が置いてあったりします。
今回は、この中から面白そうなもので、訓練済みモデルをダウンロードできるものを使って遊んでみたいと思います。

使うモデル

NVlabs の中の FUNIT というものを使って遊びたいと思います。
pix2pix のような Image-to-Image Translation のモデルのようです。
犬を別の動物に変えるような画像を生成したりしています。
訓練では動物の顔を別の動物の顔に変換するような訓練をしていますが、自分の顔の画像を入力して自分の顔を別の動物に変えてみようと思います。

準備

まずはレポジトリをクローンします。

> git clone https://github.com/NVlabs/FUNIT.git
Cloning into 'FUNIT'...
remote: Enumerating objects: 29, done.
remote: Counting objects: 100% (29/29), done.
remote: Compressing objects: 100% (27/27), done.
Remote: Total 208 (delta 0), reused 10 (delta 0), pack-reused 179R
Receiving objects: 100% (208/208), 97.38 MiB | 9.12 MiB/s, done.
Resolving deltas: 100% (85/85), done.

各種ライブラリのインストール

README.md に必要なライブラリが書いてあるのでインストールしていきます。

• install CUDA
  • https://developer.nvidia.com/cuda-downloads

• install cuDNN

  • nvidia のドライバーダウンロード

    • https://www.nvidia.com/Download/index.aspx

  • SDK をダウンロード

    • ログインが必要なのでログインします。

• アンケートを回答するとバージョン選択画面が出るので windows のものをダウンロードします。

• 解答すると以下のフォルダが作成される

• CUDA ダウンロード時に以下のフォルダが作成されているので、解答したフォルダの bin, include, lib をすべてコピーします。

• CUDA_PATH が環境変数に登録されているのを確認します。

PS C:\Users\81801> $env:CUDA_PATH
C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.2

• anaconda3 のインストール

  • scoop の extras を使用してのダウンロードができますが、レポジトリ情報がメンテされていないのでインストーラーでのダウンロードを行います
  • https://www.anaconda.com/products/individual

• ライブラリのダウンロード

  • cudatoolkit に指定するバージョンを事前に確認しておきます

> nvcc -V
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2020 NVIDIA Corporation
Built on Mon_Nov_30_19:15:10_Pacific_Standard_Time_2020
Cuda compilation tools, release 11.2, V11.2.67
Build cuda_11.2.r11.2/compiler.29373293_0

• cuda toolkit のバージョンが 11.2 であったのでライブラリダウンロード時に cudatoolkit=11.2 の指定をするべきだと思うのですが、11.2 がなかったので 11.0 をダウンロードします。

conda install -y pytorch torchvision cudatoolkit=11.1 -c pytorch
conda install -y -c anaconda pip
pip install pyyaml tensorboardX
conda install -y -c menpo opencv3

以上でインストールが終了しました。

試してみる

pretrained モデルの項目の link からチェックポイントをダウンロードし、動かしてみます。

> python test_k_shot.py --config configs/funit_animals.yaml --ckpt pretrained/animal149_gen.pt --input images/input_content.jpg --class_image_folder images/n02138411 --output images/output.jpg

README 通りの画像ができました。

自分の顔を変えてみる

次は input の画像を自分の顔にしていろんな動物を試そうと思います。
訓練元データは、似たようなモデルをまとめた 別のレポジトリ のスクリプトでダウンロードできるようなので、そちらで取得します。

ダウンロードしてきたデータは、各動物の画像が4GBほどありました。

各動物の画像を5枚選んで先程のコマンドに渡します。
input 画像の方は自分の顔を指定します。
input 画像は 128x128の画像だったので切り抜きました。

自分の顔を変換した画像の横に変換対象の動物を並べています。

わりとよさげ

メガネかけてる

毛もきれいに消えた

所感

お面を被ったような画像が結構出来上がりましたが、人間の顔なのに結構うまく変換できました。
髪の毛が結構残ったので次回はハゲにして臨もうと思います。
NVIDIAカタログには最新のGPUが大量に必要なモデルの訓練済みモデルも公開していたりするので、遊んでみると楽しいです。