Creating Abstract Art with StyleGAN2 ADA | Towards Data Science

Back in August 2020, I created a project called MachineRay that uses Nvidia’s StyleGAN2 to create new abstract artwork based on early 20th century paintings that are in the public domain. Since then, Nvidia has released a new version of their AI model, StyleGAN2 ADA, that is designed to yield better results when generating images from a limited dataset [1]. (I’m not sure why they didn't call it StyleGAN3, but I’ll refer to the new model as SG2A to save a few characters). In this article, I’ll show you how I used SG2A to create better looking abstract paintings.
Similar to the original Machine Ray, I am using abstract paintings that are in the public domain as my source images to train the SG2A system. I then change the aspect ratio as part of the post-processing. This diagram shows the flow through the system.
Everything starts with the 850 images I scraped from WikiArt.org using a custom script. The images are pre-processed and fed into the Discriminator Network as the “real” images. A set of 512 random numbers are chosen and fed into the Style Mapper and Generator Networks to create the “fake” images. Both the real and fake images are modified with Adaptive Discriminator Augmentation, the key innovation in SG2A. I’ll discuss this further down in the article. The job of the Discriminator Network is to determine if the input is real or fake. The result is fed back into the three networks to train them. When the training is done, I post-process the output of the Generator Network to get the final images.
I’ll go through each step in more detail in the sections below.
I use a Python script to gather images from WikiArt.org that are labeled as being “abstract” and are in the public domain, i.e. created before 1925. The images are from the likes of Mondrian, Kandinsky, Klee, etc. The source code to gather the images is here. Here is a sample of the source paintings.
Here are the steps I use for pre-processing the images:
Remove the frames: A Python script looks for and removes wooden (or other) frames around the images. The code is here.Resize to 1024x1024: GANs work best with square images that are resized to be a power of 2, e.g. 210 = 1024
Remove the frames: A Python script looks for and removes wooden (or other) frames around the images. The code is here.
Resize to 1024x1024: GANs work best with square images that are resized to be a power of 2, e.g. 210 = 1024
Here is a sample of pre-processed images.
One of the major improvements in SG2A is dynamically changing the amount of image augmentation during training.
Image augmentation has been around for a while. The concept is fairly simple. If you don’t have enough images to train a GAN, it can lead to poor performance, like overfitting, underfitting, or the dreaded “model collapse”, where the generator repeats the same output image. A remedy for these problems is image augmentation, where you can apply transformations like rotation, scaling, translation, color adjustments, etc., to create additional images for the training set.
A downside to image augmentation is that the transformations can “leak” into the generated images which may not be desirable. For example, if you are generating human faces, you could use 90° rotation to augment the training data, but you may not want the generated faces to be rotated. Nvidia found that augmentations can be designed to be non-leaking on the condition that they are skipped with a non-zero probability. So if most of the images being fed into the discriminator are not rotated for augmentation, the generator will learn to not create images that are rotated [1].
...the training implicitly undoes the corruptions and finds the correct distribution, as long as the corruption process is represented by an invertible transformation of probability distributions over the data space. We call such augmentation operators non-leaking. — Tero Karras, et al.
The new version of StyleGAN has a feature called Adaptive Discriminator Augmentation (ADA) that performs non-leaking image augmentations during training. A new hyperparameter, p, in the range of 0 to 1, determines how much and how often augmentations are to be applied to both the real images and the fake images during training.
Here is a sample of augmentations with different values for p.
You can see how the sample images show more variety with spatial and color changes as p increases from 0.0 to 1.0. (Note that SG2A does not show the augmentation done to the real images during training. So I added the functionality to do this in my fork of the code, here. You can run a test in the Google Colab, here.)
The value p starts at 0 when training begins, and then increases or decreases if the system senses that there is overfitting or underfitting during training.
The heuristic used in the system is based on the sign (positive or negative) of the output of the discriminator during the processing of a batch of generated images. If there are more positive output values than negative, then the trend is towards real images, which is an indication of overfitting. If there are more negative values than positive then the trend is towards fake images, which is an indication of underfitting. The value p is adjusted up or down accordingly after the batch — up for overfitting and down for underfitting.
A target value for p can be set, i.e. 0.7, so there is always a non-zero probability that the augmentations can be skipped, which avoids leaking.
You can read more about Adaptive Discriminator Augmentation in Mayank Agarwal’s post here.
I trained the system using Google Colab. It took about four days to run. The system has an Nvidia Tesla V100 GPU which can run up to 14 teraFLOPS (14 trillion floating-point operations per second). So it took about 4.6 exaFLOPS to train (about 4.6 quintillion floating-point operations). That’s a lot of math.
Here’s the shell command I used to kick things off.
python stylegan2-ada/train.py --aug=ada --target=0.7 \--mirror=1 --snap=1 --gpus=1 \--data='/content/drive/MyDrive/datasets/paintings' \--outdir='/content/drive/MyDrive/results_stylegan2_ada'
Setting the aug parameter to ada enables the Adaptive Discriminator Augmentation to kick in when needed. Setting the target parameter to 0.7 prevents p from going over 0.7, which maximizes the dynamic augmentation without leaking any further augmentation into the final images.
Here’s a graph of the training results over time. I am using the Fréchet Inception Distance (FID) as a metric for image quality and diversity, where a lower score is better [2]. You can read about the FID score in Cecelia Shao’s post here.
You can see how StyleGAN2 ADA outperforms the original StyleGAN2 for the same number of iterations. The FID score for SG2A bottomed out at just over 100 after about 300 iterations. Here are some samples of the results.
The GAN seems to produce a nice variety of abstract paintings with interesting compositions with varied color schemes. But there is room for improvement.
It may not be intuitive, but it’s possible to improve the quality of the paintings by first training the GAN on a different, larger set of images, and then further train the model using the abstract paintings. This technique is called Transfer Learning. It was first described as a technique for training NNs in 1976 [3].
Transfer Learning is technique that uses a pre-trained neural network trained for Task 1 for achieving shorter training time in learning Task 2. [4] — Stevo Bozinovski
I remember learning in my Art History 101 class that many abstract painters started by painting figurative subjects, like people and landscapes. For example here are sequences of paintings by Mondrian, Klee, and Kandinsky that show their progressions from landscape to abstract art.
I wanted to see if I could transfer the learning from creating landscape photos to perhaps improve its ability to create abstract paintings.
As an experiment, I trained SG2A on 4,800 photographs of landscape photos from Flickr that are in the public domain. The dataset, by Arnaud Rougetet, is available on Kaggle, here.
Here is a sample of the photos from Flickr. Note that I resized the photos to be 1,024 x 1,024 pixels each.
I trained the SG2A system using another Google Colab to create new landscape images. It also took about four days to run the training. Here’s the shell command I used.
python stylegan2-ada/train.py --aug=ada --target=0.7 \--mirror=1 --snap=1 --gpus=1 \--data='/content/drive/MyDrive/datasets/landscapes' \--outdir='/content/drive/MyDrive/results_stylegan2_ada_landscapes'
The command is almost the same as the one I used to train using the abstract paintings. The only differences are the paths to the landscape folders on Google Drive. Here’s a graph that shows the FID scores of the training of landscapes (the green line).
The score for the generated landscapes is much better than the abstract paintings. You can see that it bottomed out around 25 after about 150 iterations. The reason the landscapes have a better score than the abstract paintings is probably due to the larger training set — 4,800 landscapes and only 850 abstract paintings.
Here are some generated landscape photos that were created after four days of training SG2A.
These look pretty good. The horizons are a bit distorted and the clouds seem somewhat fantastical, but they could otherwise pass as real landscape photos. Now let’s see if this learning can help with creating abstract art.
The next step of the experiment was retraining SG2A to create abstract art starting with the model previously trained on landscapes. Here’s a flow diagram that shows the process.
After the GAN was trained to create landscapes, it was further trained to create abstract paintings using prior knowledge of the landscapes. I trained the system using a third Google Colab. Here’s the shell command I used.
python stylegan2-ada/train.py --aug=ada --target=0.7 \--mirror=1 --snap=1 --gpus=1 \--data='/content/drive/MyDrive/datasets/paintings' \--resume='/content/drive/MyDrive/results_stylegan2_ada/00006-landscapes-mirror-auto1-ada-target0.7-resumecustom/network-snapshot-000048.pkl' \--outdir='/content/drive/MyDrive/results_stylegan2_ada_tl'
I am training using abstract paintings, but the resume option starts the training with the GAN trained with landscapes. It took about four days to run the final training. Here is a graph of the FID scores. The new results are in blue.
You can see that Transfer Learning helped improve the scores. The FID numbers for the latest training run settle in around 85, which is better than the previous run with a score of about 105. Here are some sample images that benefitted from the Transfer Learning.
It may be subtle, but the images created with the model with Transfer Learning appear to be more refined, pleasing compositions. Please check out the appendix below for a gallery of samples.
Once the GAN is trained, images can be generated using this command line.
python stylegan2-ada/generate.py --seeds=1-4 --trunc=1.5 \  --outdir results --network=network-snapshot-000188.pkl
This will generate four images using random seeds 1, 2, 3, and 4. The truncation parameter, trunc, will determine how much variation there will be in the images. I found that the default of 0.5 is too low, and 1.5 gives much more variety.
A Google Colab for generating images with a variety of aspect ratios is available here. It generates 21 images and lets you choose one to see at high resolution. Here is a sample of images.
Additional work might include training on paintings of landscapes instead of photos for the Transfer Learning. This may help the GAN pick up on a painterly look to apply to abstract paintings.
I would like to thank Jennifer Lim, Oliver Strimpel, Mahsa Mesgaran, and Vahid Khorasani for their help and feedback on this project.
The 850 abstract paintings I collected can be found on Kaggle here. All source code for this project is available on GitHub. The sources are released under the CC BY-NC-SA license.
[1] Karras, T., Aittala, M., Hellsten, J., Laine, S., Lehtinen, J., and Aila., T., “Training Generative Adversarial Networks with Limited Data.”, October 7, 2020, https://arxiv.org/pdf/2006.06676.pdf
[2] Eiter, T. and Mannila, H., “Computing the Discrete Fréchet Distance”, Christian Doppler Labor für Expertensyteme, April 25, 1994, http://www.kr.tuwien.ac.at/staff/eiter/et-archive/cdtr9464.pdf
[3] Bozinovskim S., and Fulgosi, A., “The influence of pattern similarity and transfer learning upon training of a base perceptron B2.” Proceedings of Symposium Informatica, 3–121–5, 1976
[4] Bozinovski, S., “Reminder of the first paper on transfer learning in neural networks, 1976”. Informatica 44: 291–302, 2020, http://www.informatica.si/index.php/informatica/article/view/2828
Some rights reserved