StyleGAN2 on Jetson Nano

Category:

AI

Tags:

AIDeep LearningJetson NanoNVIDIAStyleGAN

Last Updated: 2022/02/11 10:21:40

I wrote a post about StyleGAN on Jetson Nano, and then I will note about also StyleGAN2.

＊The basic environment construction is the same as the following article.

• Set Up Jetson Nano Developer Kit B01
• StyleGAN on Jetson Nano

Configure $PATH for nvcc

First, I'm now using JetPack 4.6 but it doesn't have a path to nvcc command in CUDA, so I had to configure the path of /usr/local/cuda/bin (where nvcc is located) into the $PATH variable.

# You need to configure the path to `nvcc`

# Check if the path is already configured for nvcc
# If nvcc can be found, the following configuration is not necessary.
$ which nvcc

# If nvcc can't be found, we need to configure the PATH for nvcc.
# nvcc is located in `/usr/local/cuda/bin`
$ echo 'export PATH=${PATH}:/usr/local/cuda/bin' >> ~/.bashrc
$ export PATH=${PATH}:/usr/local/cuda/bin

Then, clone StyleGAN2 from GitHub.

StyleGAN2 (opens new window)

Requirements in README.md says SytleGAN2 uses nvcc so first you need to make sure it works correctly. You can test it with test_nvcc.cu.

$ cd ~
$ git clone https://github.com/NVlabs/stylegan2
$ cd stylegan2

# You can test if nvcc works well with the following command.
# 
# Successful if the following lines are displayed.
# 
# CPU says hello.
# GPU says hello.

$ nvcc test_nvcc.cu -o test_nvcc -run

Execute Image Generations

Then, as following the instruction written in README.md, execute image generations.

# generate images from an already trained model.
# but in the case of Jetson Nano, default python command is python2.7 and (at lease if you are following my posts) now you can use python3 command.
python3 run_generator.py generate-images --network=gdrive:networks/stylegan2-ffhq-config-f.pkl --seeds=6600-6625 --truncation-psi=0.5

Now then the following error occurred 😦

tensorflow.python.framework.errors_impl.NotFoundError: /home/kaz/stylegan2/dnnlib/tflib/_cudacache/fused_bias_act_601777699de14239ee94003371c74f59.so: undefined symbol: _ZN10tensorflow12OpDefBuilder5InputESs

I googled this error back and forward and I found this is caused from TensorFlow version differences, but I couldn't fix this by re-installing TensorFlow and other related libraries with various versions. The version of TensorFlow was originally 1.15.

$ python3 -c 'import tensorflow as tf; print(tf.__version__)' 
# 1.15.*

After some trial and error, finally I could fix the problem by modifying dnnlib/tflib/custom_ops.py with referring to this page (opens new window)(Japanese).

--- dnnlib/tflib/custom_ops.py.bk	2022-02-02 09:49:37.683383645 +0900
+++ dnnlib/tflib/custom_ops.py	2022-02-02 10:40:29.145091681 +0900
@@ -124,7 +124,7 @@
             compile_opts += '"%s"' % os.path.join(tf.sysconfig.get_lib(), 'python', '_pywrap_tensorflow_internal.lib')
         elif os.name == 'posix':
             compile_opts += '"%s"' % os.path.join(tf.sysconfig.get_lib(), 'python', '_pywrap_tensorflow_internal.so')
-            compile_opts += ' --compiler-options \'-fPIC -D_GLIBCXX_USE_CXX11_ABI=0\''
+            compile_opts += ' --compiler-options \'-fPIC -D_GLIBCXX_USE_CXX11_ABI=1\''
         else:
             assert False # not Windows or Linux, w00t?
         compile_opts += ' --gpu-architecture=%s' % _get_cuda_gpu_arch_string()

Now Try again, and it doesn't work well again because of memory shortage 😦

I added the following lines to run_generater.py as the same as trying StyleGAN. (add lines -> get killed -> remove）

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)

Try again, the results directory was created and then images were created 😃

python3 run_generator.py generate-images --network=gdrive:networks/stylegan2-ffhq-config-f.pkl --seeds=6600-6625 --truncation-psi=0.5

Other Examples of Pre-Trained Networks

There are 4 examples of pre-trained networks in README.md, and the above command if the first item of them.

The 2nd one generates curated images.

python3 run_generator.py generate-images --network=gdrive:networks/stylegan2-ffhq-config-f.pkl  --seeds=66,230,389,1518 --truncation-psi=1.0

The 3rd one generates car images.

python3 run_generator.py generate-images --network=gdrive:networks/stylegan2-car-config-f.pkl  --seeds=6000-6025 --truncation-psi=0.5

The last one generates images by mixing styles, but when I tried this with Jetson Nano, it took more than 10 hours and terminated before image generation was complete. It seems so heavy 😦