[PYTHON] Try StyleGAN on Google Colaboratory

Overview

When I tried StyleGAN on Google Colaboratory, I had some stumbling points, so I will summarize it as a memorandum.

The reference Web pages are as follows. Once StyleGAN works properly, you may want to refer to these well-explained pages.

-"I tried using StyleGAN with Google Colaboratory. ” -"Python implementation of DCGAN, CycleGAN & StyleGAN" -"StyleGAN" The era when photographs are evidence is over. ""

The points I stumbled upon are as follows.

--Connect to the GPU ** with Google Colaboratory . - The version of tensorflow must be ** to 1.14 or 1.15. -** tensorflow-gpu must be prepared at the same time **.

1. Preparation

Select GPU for Google Colaboratory runtime

Start Google Colaboratory, click "Runtime"> "Change Runtime Type" from the toolbar at the top of the screen, specify "GPU" in the hardware accelerator, and click "Save". This will switch the runtime type from CPU to GPU.

Version of TensorFlow

StyleGAN works with a slightly older version of tensorflow, so you have to downgrade from the latest tenforflow. As of November 5, 2020, the latest version of tensorflow installed in Google Colaboratory is 2.3.0. The confirmation method is as follows.

import tensorflow as tf
print(tf.__version__)

output

2.3.0

According to this page

Also generally (NVlabs) Stylegan and Stylegan2 require TensorFlow 1.14 or 1.15 with GPU support.

Therefore, it is necessary to prepare tensorflow 1.14 or 1.15. The following code installs tensorflow 1.15.

!pip install tensorflow==1.15

GPU support for TensorFlow

In versions prior to 1.15, the TensorFlow CPU package and GPU package are separate. (-> "GPU Support | TensorFlow") Therefore, it is necessary to prepare tensorflow that can use GPU. This can be solved by specifying tensorflow-gpu.

!pip install tensorflow-gpu==1.15

At this point, restart the runtime once. After that, when the reconnection is completed, execute the following code, and if you get the same output, the GPU specification and TensorFlow version and package should be prepared correctly.

import tensorflow as tf
print(tf.__version__)
print(tf.test.gpu_device_name())

output

1.15.0
/device:GPU:0

Cooperation with Google Drive

Let's use Google Drive as a storage place for related files.

from google.colab import drive
drive.mount('/content/drive')

Create a directory to place StyeGAN related files in advance and move it there.

%cd /content/drive/My Drive
!mkdir stylegans
%cd stylegans
!mkdir stylegans_dir

Download StyleGAN related files

!git clone https://github.com/NVlabs/stylegan.git
!git clone https://github.com/NVlabs/stylegan2.git

The distributed trained model can be obtained by specifying the URL, but there seems to be a problem that the runtime ends. To avoid this, download only the required trained model separately.

Since the model karras2019stylegan-ffhq-1024x1024.pkl is placed in this page, after downloading it to your own PC, Google Drive Save it directly under the stylegans directory of.

2. Load the trained Generator

%cd /content/drive/My Drive/stylegans/stylegan
 
import os
import pickle
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
import dnnlib
import dnnlib.tflib as tflib
import config
 
# Initialize TensorFlow.
tflib.init_tf()
 
# Load pre-trained network.
*_, Gs = pickle.load(open('../karras2019stylegan-ffhq-1024x1024.pkl','rb'))

3. Definition of latent variables

# Pick latent vector.
rnd = np.random.RandomState(210)
latents = rnd.randn(1, Gs.input_shape[1])

If you change 210 in RandomState (210) to another number, various types of images will be displayed.

4. Generate and display an image

# Generate image.
fmt = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True)
images = Gs.run(latents, None, truncation_psi=0.7, randomize_noise=True, output_transform=fmt)
plt.imshow(images[0])
plt.tick_params(labelbottom="off",bottom="off")
plt.tick_params(labelleft="off",left="off")

If you run the above code, you should see an image like the one below.

5. Change the parameters and play with

Latent variable

The images generated by changing the numerical value specified by RandomState () are displayed side by side.

plt.figure(figsize=(15,10))

for i in range(20):
  rnd = np.random.RandomState(i*20+10)
  latents = rnd.randn(1, Gs.input_shape[1])
  
  plt.subplot(4,5,i+1)
  fmt = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True)
  images = Gs.run(latents, None, truncation_psi=0.7, randomize_noise=True, output_transform=fmt)
  plt.title("RandomState(" + str(i*20+10) + ")")
  plt.imshow(images[0])
  plt.tick_params(labelbottom="off",bottom="off")
  plt.tick_params(labelleft="off",left="off")
  plt.axis('off')

By changing the latent variable, it can be confirmed that completely different types of images (men and women, adults and children, etc.) are output.

I tried to generate more images. You can see that there are considerable variations.

truncation_psi

I changed truncation_psi, which is specified as an argument ofGs.run ().

plt.figure(figsize=(30,10))

for i in range(11):
  rnd = np.random.RandomState(210)
  latents = rnd.randn(1, Gs.input_shape[1])
  
  plt.subplot(1,11,i+1)
  fmt = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True)
  temp = -1 + 0.2*i
  images = Gs.run(latents, None, truncation_psi=temp, randomize_noise=True, output_transform=fmt)
  plt.title(str(round(-1 + 0.2*i,1)))
  plt.imshow(images[0])
  plt.tick_params(labelbottom="off",bottom="off")
  plt.tick_params(labelleft="off",left="off")
  plt.axis('off')

I put two images, but if you look closely, the details such as hair and background pattern are slightly different. This seems to be because randomize_noise = True contributes.

For RandomState (290)

in conclusion

This page summarizes the steps to get StyleGAN working with Google Colaboratory. For more information about StyleGAN and other ways to use it, there are already many articles and web pages, so please refer to them.