Jetson NanoでStyleGANを試してみる
Jetson Nano で StyleGAN を試してみます。
マシンスペック的に学習モデルの生成は難しいと思われるのですが、学習済みのモデルを用いての画像生成は普通にできたので、その手順をまとめます。
https://github.com/NVlabs/stylegan (opens new window)
環境構築はからあげさんという方が作成した下記の2つの記事が役にたちました。
Jetson Nano上でStyleGANを動かすだけなら基本的にこれだけでも足りるはずかと思います。
- Jetson Nano関係のTIPSまとめ (opens new window)
- Jetson NanoにJetPack 4.4/4.5を入れてTensorFlow・物体検出・姿勢推定・ROS2(Realsense)・ROS1を動かす (opens new window)
# 前準備
まずはJetPack初期化済みのJetson NanoにSSHログインします。
（Jetson Nanoの初期化方法はこちらを参照）
$ ssh kaz@192.168.55.1
スワップ確保。前回はチュートリアルに沿ってコマンドでスワップファイルを作成しましたが、githubにスワップファイルの生成スクリプトが公開されておりましたのでありがたく利用させていただくことにします 🙏
$ git clone https://github.com/JetsonHacksNano/installSwapfile
$ cd installSwapfile
$ ./installSwapfile.sh
# check
$ free -m
2
3
4
5
6
前述のからあげさんがJetPack上で機械学習に必要なライブラリをまとめてダウンロードするためのスクリプト (opens new window)を公開しているので、こちらもありがたく利用させていただきます 🙏
私自身はJetPack 4.6で試しているのに対してインストールスクリプトはJetPack4.5までの記述しかありませんが、JetPack4.5用のスクリプトで特に問題はなさそうでした。ディープラーニング関連あるあるのようですが、最初はあちこちのサイトを参考に必要なライブラリを落としていたものの依存関係の問題なのかうまくいかなかったので、Jetson Nanoでテストする場合、今のところはこちらのスクリプトで必要なライブラリをまとめてインストールしてしまうのが楽だと思います。
$ git clone https://github.com/karaage0703/jetson-nano-tools
$ cd jetson-nano-tools
$ ./install-tensorflow-v45.sh
2
3
さらに上記のスクリプトに含まれませんが、StyleGANを動かすにはPillowが必要なので別途、インストールします。
$ sudo pip3 install -U Pillow
# StyleGANのダウンロードと実行
githubからstyleganをクローンして、pretrained_example.py という、NVIDIAが用意した学習モデルを元に画像を生成するサンプルプログラムを試してみます。
$ cd ~
$ git clone https://github.com/NVlabs/stylegan
$ cd stylegan
$ python3 pretrained_example.py
2
3
4
なんかGoogle Driveから学習モデルがダウンロードできない（OSError: Google Drive quota exceeded）とかって怒られたので、Google Drive上のファイル (opens new window) に直接アクセスしてダウンロードします。
無事、ダウンロードできたらファイルをJetson Nanoに転送します。
（Jetson Nano側のブラウザで直接ダウンロードした場合は不要）
# Jetson Nano側で直接、ダウンロードした場合は不要
# ユーザー名などは個々の状況に合わせて修正してください
$ scp ~/Downloads/karras2019stylegan-ffhq-1024x1024.pkl kaz@192.168.55.1:/home/kaz/stylegan/
2
3
さらにpretrained_example.py についてGoogle Drive読み込み部分を下記のようにローカルファイルの読み込みに修正し、再度、スクリプトを実行します。
--- pretrained_example.py.bk 2022-02-01 15:29:21.141652160 +0900
+++ pretrained_example.py 2022-02-01 15:44:22.893303912 +0900
@@ -20,8 +20,9 @@
tflib.init_tf()
# Load pre-trained network.
- url = 'https://drive.google.com/uc?id=1MEGjdvVpUsu1jB4zrXZN7Y4kBBOzizDQ' # karras2019stylegan-ffhq-1024x1024.pkl
- with dnnlib.util.open_url(url, cache_dir=config.cache_dir) as f:
+ # url = 'https://drive.google.com/uc?id=1MEGjdvVpUsu1jB4zrXZN7Y4kBBOzizDQ' # karras2019stylegan-ffhq-1024x1024.pkl
+ # with dnnlib.util.open_url(url, cache_dir=config.cache_dir) as f:
+ with open("./karras2019stylegan-ffhq-1024x1024.pkl" , "rb") as f:
_G, _D, Gs = pickle.load(f)
# _G = Instantaneous snapshot of the generator. Mainly useful for resuming a previous training run.
# _D = Instantaneous snapshot of the discriminator. Mainly useful for resuming a previous training run.
2
3
4
5
6
7
8
9
10
11
12
13
14
再度実行したところ、「Resource exhausted: OOM when allocating tenso」とかいう謎のエラーがでました。。。
調べてみたところ、メモリが足りなくてエラーが発生しているらしいので、ここ (opens new window) を参考にプロセスでTensorFlowが使うメモリを制限するようにして再度、実行します。
リンク先のサイトに書いてあるものと同じですが、pretrained_example.py の冒頭に下記のように追加して実行すると、プロセスがkillされるので、そしたら追加した箇所をコメントアウトして再度、実行します。
import tensorflow as tf
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.7)
sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))
config = tf.ConfigProto()
config.gpu_options.allow_growth=True
sess = tf.Session(config=config)
2
3
4
5
6
7
resultsフォルダが生成され、その中に下記のようなおじさんの顔が生成されれば成功です 😃
ちなみにStyleGANでは画像生成の際に潜在変数と呼ばれる変数を渡していますが、それを変えると異なる画像となります。
下記は（見づらいのですが）少しづつ潜在変数を変えて画像を生成した場合の表示です。
