Generating your own Images with NVIDIA StyleGAN2-ADA for PyTorch on the Ampere Architecture | by Jeff Heaton | Towards Data Science

StyleGAN2 ADA allows you to train a neural network to generate high-resolution images based on a training set of images. The most classic example of this is the made-up faces that StyleGAN2 is often used to generate. Until the latest release, in February 2021, you had to install an old 1.x version of TensorFlow and utilize CUDA 10. This requirement made it difficult to leverage StyleGAN2 ADA on the latest Ampere-based GPUs from NVIDIA. In this article I will show you how to use this new version of StyleGAN from Windows, no Docker or Windows Subsystem for Linux (WSL2) needed!
I’ve trained GANs to produce a variety of different image types, you can see samples from some of my GANs above. I provide pretrained models to produce these images on GitHub:
Minecraft GAN
70s Scifi Art GAN
Fish GAN
Christmas GAN
You can make use of the above networks, using only Google Colab online, to generate these sorts of images for yourself.
GANs are compute-intensive, there is really no way around it. The researchers at NVIDIA threw 8 V100s on a DGX system at the task of training faces. StyleGAN also scales nearly linearly across multiple GPUs, as a result, StyleGAN will take whatever hardware you are willing to throw at it on a single machine. The images that I present here were trained on a dual Quadro RTX 8000.
For this article, I am assuming that we will use the latest CUDA 11, with PyTorch 1.7.1. NVIDIA recommends 12GB of RAM on the GPU; however, it is possible to work with less, if you use lower resolutions, such as 256x256. Higher-resolution GANs are generally trained at 1024x1024. A single strong GPU, such as an NVIDIA RTX A6000 also does very well.
We will install StyleGAN2 outside of WSL2 or Docker. This gives the highest performance. Believe me, with GANs, you want every bit of compute that your machine can provide!
First, make sure you have the latest NVIDIA driver for your graphics card:
NVIDIA GPU Drivers
Second, install the latest version of CUDA 11. If CUDA 12 has been released, and I have not updated this article, then proceed with caution. Check the StyleGAN2 ADA PyTorch instructions for the latest updates on versions.
CUDA Toolkit Downloads
It is also necessary that you install Visual C++ so that custom CUDA kernels can be compiled by StyleGAN. Visual Studio Community edition can be found at the following URL. Make sure that you install C++, which is not enabled by default in Microsoft’s installer.
https://visualstudio.microsoft.com/vs/
Once Visual Studio is installed, you must add several items to the system PATH and environmental variables. You can accomplish this by running the following batch program:
C:\Program Files (x86)\Microsoft Visual Studio\<VERSION>\Community\VC\Auxiliary\Build\vcvars64.bat
I suggest opening a power shell window and moving into that directory and running the command “vcvars64.”
To install PyTorch you will need Python installed on your system. I suggest either installing Miniconda or Anaconda. Miniconda is small and you will have to install needed packages. Anaconda is gigantic but contains mostly everything you will ever need.
https://docs.anaconda.com/anaconda/install/
https://docs.conda.io/en/latest/miniconda.html
If you are even debating which to install, or you’ve never heard of the difference between Anaconda and Miniconda, then almost certainly, Anaconda is the right answer for you.
The StyleGAN team recommends PyTorch 1.7.1 for StyleGAN. Later versions may likely work, depending on the amount of “breaking changes” introduced to PyTorch. The install process for PyTorch is amazing, navigate to the following URL and choose your options:
https://pytorch.org/
I choose:
Stable 1.7.1
Windows
PIP
Python
CUDA 11
This resulted in the following command:
pip install torch===1.7.1+cu110 torchvision===0.8.2+cu110 torchaudio===0.7.2 -f https://download.pytorch.org/whl/torch_stable.html
The actual NVIDIA StyleGAN2 ADA package is distributed through GitHub at the following repository:
https://github.com/NVlabs/stylegan2-ada-pytorch
You can choose to download a ZIP file, which should be extracted to a directory. You can also obtain StyleGAN with the command line git command.
git clone https://github.com/NVlabs/stylegan2-ada-pytorch.git
At the time of this writing, there was not a requirements.txt file provided by StyleGAN to specify required packages. As you run the Python scripts, you will see errors about missing packages, just pip install them. To get you started, these are the ones I found:
pip install clickpip install tqdmpip install requestspip installimageiopip installpsutilpip installscipy
To train the GAN you must convert all of the images to PNG of the same size and dimension, with a very specific directory structure. This can be accomplished with the dataset_tool script provided by StyleGAN. Here I am converting all of the JPEG images that I obtained to train a GAN to generate images of fish.
python dataset_tool.py --source c:\jth\fish_img --dest c:\jth\fish_train
Next, you will actually train the GAN. This is done with the following command:
python train.py --data c:\jth\fish_train --outdir c:\jth\results
Training can take a long time, as it trains, you will see images created during training that indicate how good the images currently look. The following image shows my Minecraft GAN at 2,800 kimg, meaning it had trained on over 2,800 images (both real and augmented), at this point.
As training progresses pickle snapshots will be made of your generator and discriminator to correspond with each image set checkpoint generated. For example, network-snapshot-002800.pkl generated the image checkpoint above. To generate images from this network, the following command is used.
python generate.py --outdir=out --trunc=1 --seeds=85,265,297,849 \    --network=network-snapshot-002800.pkl
The seeds specify the individual images to be generated, each seed is a separate image.
If you would like to see each of these steps performed, I also have a YouTube video of this same material.
NVIDIA StyleGAN2 ADA is a great way to generate your own images if you have the hardware for training. The new PyTorch version makes it easy to run under a Windows environment. Following the steps in this article allows you to quickly setup an environment for training your own GANs.