Merry GAN-mas: Introduction to NVIDIA StyleGAN2 ADA | by Jeff Heaton | Towards Data Science

Generative Adversarial Neural Networks (GANs) are a type of neural network that can generate random “fake” images based on a training set of real images. GANs were introduced by Ian Goodfellow in his 2014 paper. GANs trained to produce human faces have received much media attention since the release of NVIDIA StyleGAN in 2018. Websites like Which Face is Real and This Person Does Not Exist demonstrate the amazing capabilities of NVIDIA StyleGAN. In this article I will explore the latest GAN technology, NVIDIA StyleGAN2 and demonstrate how to train it to produce holiday images.
The first step is to obtain a set of images to train the GAN. I created a Python utility called pyimgdata that you can use to download images from Flickr and perform other preprocessing. Flickr is a great place to obtain images and is used by many GAN paper authors, such as NVIDIA. Flickr is beneficial because it has an API to obtain images and contains license information for each upload. When building a dataset of images, it is generally advisable to use only images published by their authors with a permissive license.
My Flickr download utility makes use of a configuration file, such as the following:
This script downloads the results of the specified search into the specified path. The filenames will have the specified prefix. I specify all licenses because I do not intend to publish this image list. This actually brings up an open issue in copyright law. If a neural network learns from copyrighted and produces new work, is the AI bound to the original copyright? Similarly, is a human musician who listens to copyrighted music beholden to the copyright owner for inspiration that the music had on the musician's brain? For the purposes of copyright, I consider my GAN and its images to be a derivative work.
The license specifies which licenses you wish to download:
0 — All Rights Reserved
1 — Attribution-NonCommercial-ShareAlike License
2 — Attribution-NonCommercial License
3 — Attribution-NonCommercial-NoDerivs License
4 — Attribution License
5 — Attribution-ShareAlike License
6 — Attribution-NoDerivs License
7 — No known copyright restrictions
8 — United States Government Work
9 — Public Domain Dedication (CC0)
10 — Public Domain Mark
The process section of the configuration file allows the program to perform basic preprocessing on the images. For StyleGAN 2, the images must be square.
For this project, I collected around 3,000 images from Flickr, a sampling of which you can see here:
Training a GAN requires considerable computational power, particularly GPUs. Even with multiple high-end GPUs, training will take days. For this project, I had access to a Lenovo ThinkPad P53 and ThinkStation P920 that Lenovo had loaned me for my YouTube channel. Having access to dual high-end GPUs considerably speeds the training process.
StyleGAN2 ADA requires TensorFlow 1.14, which is somewhat older than the TensorFlow 2.x that I currently run on my machine. Additionally, TensorFlow 1.14 would require older CUDA drivers than I presently have installed. Because of these reasons, I created a Docker image to run StyleGAN2 ADA.
This Docker image can be run with the following command. This assumes you have NVIDIA Docker installed.
nvidia-docker run -it -u $(id -u):$(id -g) -v /home/username/:/mnt --device /dev/nvidia0:/dev/nvidia0 --device /dev/nvidiactl:/dev/nvidiactl --device /dev/nvidia-uvm:/dev/nvidia-uvm stylegan2:1.0 /bin/bash 
The above command establishes the following:
-u $(id -u):$(id -g) — This causes the user to be logged in to Docker to be the same as the user running the Docker command. This ensures that all permissions and ownerships are correct on the mounted volumes.
/home/username/:/mnt — You should mount a volume to access your images and store results.
/bin/bash — We start the BASH shell to allow you to execute commands.
Your images must be converted to TFRecords. All of your images must be of the same size, square, and color depth. You cannot mix color and greyscale images. The following command converts a data set named “Christmas”. The source JPEGs should be in the /mnt/data/christmas, and the resulting TFRecords will be written to /mnt/datasets/christmas.
python dataset_tool.py create_from_images /mnt/datasets/christmas /mnt/data/christmas
To actually train your GAN, you will enter a command similar to the following.
python train.py --gpus=1 --data=/mnt/datasets/christmas --mirror=1 --kimg 3000 --outdir=/mnt/results --aug=ada
You must provide an input directory that contains your image TFRecords. You must also provide an output results directory. The results directory will contain your generator's snapshots, sample images as training progresses, and a log.
While it takes considerable processing power to train a GAN, producing images with a GAN takes considerably less processing power. A GPU is still required to generate images; however, the GPU provided by Google CoLab is sufficient to perform this training. Images generated by the GAN will appear as follows:
The complete Jupyter notebook to generate your own Christmas images is stored here.
StyleGAN2 images are generated with a 512 number latent vector. The GAN was trained to generate images for random input vectors. Such a large vector is tedious to enter into source code, so seed values are generally used to generate this vector. While two very similar latent vectors will produce very similar looking images; two seeds with close numeric values will not necessarily produce similar images. Seeds are typically integer values, such as 1001, 1002, etc. The following function expands a seed to the full 512-number vector for a GAN generator specified by Gs.
def seed2vec(Gs, seed):  rnd = np.random.RandomState(seed)  return rnd.randn(1, *Gs.input_shape[1:])
The GAN contains several random values that can be changed for a given latent vector or seed. Changing these values will not substantially change a particular vector’s output image. Generally, changing this value might change hair placement or some other minor, aesthetic appearance for a face-GAN. You need to provide a second seed is provided to initialize these random values.
def init_random_state(Gs, seed):  rnd = np.random.RandomState(seed)   noise_vars = [var for name, var in Gs.components.synthesis.vars.items() if name.startswith('noise')]  tflib.set_vars({var: rnd.randn(*var.shape.as_list()) for var in noise_vars}) # [height, width]
Finally, we are ready to generate the image. The return value from this function is a numpy RGB array for the generated image.
def generate_image(Gs, z, truncation_psi):    # Render images for dlatents initialized from random seeds.    Gs_kwargs = {        'output_transform': dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True),        'randomize_noise': False    }    if truncation_psi is not None:        Gs_kwargs['truncation_psi'] = truncation_psi    label = np.zeros([1] + Gs.input_shapes[1][1:])    images = Gs.run(z, label, **Gs_kwargs) # [minibatch, height, width, channel]    return images[0]
The provided GitHub notebook also demonstrates how to fine-tune a GAN and generate a transformation video. Both of these features relate to the direct modification of the latent vector.
I produced a YouTube video covering much of the same topics as this article to see the transformation videos and fine-tuning.