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![4d86c00f-e566-4ce0-bc01-79d8300f7e9b.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/164275/e5d6497c-9bd1-8dc4-5cec-02d4933e78f7.png)
- 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 通りの画像ができました。