StyleGAN v2: notes on training and latent space exploration | by 5agado | Towards Data Science

What: in this entry I’m going to present notes, thoughts and experimental results I collected while training multiple StyleGAN models and exploring the learned latent space.
Why: this is a dump of ideas/considerations that can range from the obvious to the “ holy moly”, meant to possibly provide insight — or starting point for discussions — for other people out there with similar interests and intents. As such it can be skimmed through to see if anything is of interest.
I’m also here heavily leveraging Cunningham’s Law.
Who: many considerations apply to both StyleGAN v1 and v2, but all generated results are from v2 models, unless explicitly specified.
In term of code resources:
StyleGAN v1 and v2 official repos
Encoder for v1 (+ directions learning) and encoder for v2
All of the content here has been generated on top of such repos via my customized Jupyter notebooks.
Worth mentioning people who diligently shared code, experiments and valuable suggestions: Gwern (TWDNE), pbaylies, gpt2ent, xsteenbrugge, Veqtor, Norod78, roadrunner01. All more than worth following.
For a custom dataset, a set of images needs to be created beforehand making sure that all entries have the same square-shape and color space. You can then generate the tf-records as explained here. As it has been pointed out many times, the model is very data-hungry. Dataset size requirements vary greatly based on image content, but in general aim at 4 to 5 orders of magnitude, say around 50k with mirroring and double without. This applies to training from scratch, while in case of fine-tuning there have been many experiments done on very small datasets, providing interesting results.
I experimented mainly with two “fashion” datasets
dresses (~30k packshot images)
footwear (~140k packshot images)
But will also showcase some experiments on manually curated artistic dataset (~1k images each).
Notice the higher content consistency of the fashion datasets (in terms of location, background, etc.). compared to the artistic entries. As such, irregardless of dataset size, we expect to be able to reach higher generation accuracy for the former.
Dresses images are a case where mirror augment should be enabled, as we can double our training set for free, learning the semantic invariant property of dresses for horizontal mirroring.
Footwear is instead a case where mirroring should be disabled, as all images in the dataset are aligned, and learning mirrored versions would just be an unnecessary burden for the network.
Regarding size, I opted for 512 to first validate the results. From previous experiments with other networks I believe this is a good resolution for capturing patterns, material texture and small details like buttons and zippers. I also plan a future iteration on 256, and validate the quality of the learned representations and comparison with the training regime. I suggest moving to full 1024 resolution (or above) just if you have the proper hardware resources and have validated your dataset already against smaller resolutions.
One can train a model from scratch, or piggyback on previous trained results, fine-tuning the learned representation to a new dataset.
Training from scratch is as simple as running the following
python run_training.py -num-gpus=2 -data-dir=<your_data_dir> -config=config-f -dataset=<your_dataset_name> -mirror-augment=True -result-dir=<your_results_dir>
We already talked about the mirror-augment and the number of GPUs simply depends on your setup.
Configurations are used to control the architecture and properties of the model, starting from config-a (the original SyleGAN v1), and incrementally adding/changing components (weight demodulation, lazy regularization, path length regularization) to config-e (the full updated setup for v2), plus config-f, a larger network version of config-e and the default used in the paper.I opted on config-e for my 512 size datasets. The complexity added in config-f was mostly motivated by the information bottleneck noticed while exploring the implicit focus of the network on higher resolutions while training, so my rationale was that the added complexity would just be more detrimental in terms of training stability.
Fine-tuning is used to save time, relying on the structure already learned by a previous model, and trying to adapt such structure to a new target dataset. I have also seen people fine-tuning using images not entirely related to the original dataset, so that’s a lot up to your goals. Fine-tuning gives quicker visual feedback on training quality, so it is easier to run multiple iterations and experiments, making it a great way to initially learn more about model behavior and hyperparameters.
To fine-tune an existing network one needs to specify the snapshot to restart from. For v2 this is done in run_training.py, by adding and adapting the following
train.resume_pkl = "your_network_snapshot.pkl"  # Network pickle to resume training from, None = train from scratch.train.resume_kimg = 0.0  # Assumed training progress at the beginning. Affects reporting and training schedule.train.resume_time = 0.0  # Assumed wallclock time at the beginning. Affects reporting.
Notice how it is important to specify also the resume_kimg in accordance with your snapshot, as this is used to control the training schedule (e.g. learning rate decay).
While training the primary metric I am looking at is the Frechet Inception Distance (FID), which by default is computed every 50 training ticks on 50k images. Other metrics can be called as described in the repo README. Notice that Tensorflow logs are generated in the target destination dir and can be accessed by running tensorboard
tensorboard --logdir <LOGS_TARGET_DIR>
Apart from FID you can here inspect scores/fake and scores/real, which is the discriminator prediction for fake and real images respectively. As the label can be misleading at first, it is worth pointing out again that these are not loss values, but the pure prediction output of the discriminator.
I try to gain some guidance from such graphs, but find it often hard to extract precise insight on how the network is doing. Ideally one would like the discriminator to be unable to distinguish fake from real, and as such tending toward zero, as it is just random guessing.Other secondary metrics are present, also based on the loss choice. For example, by default the discriminator uses logistic loss with R1 regularizer, providing a plot for the gradient penalty.
For my datasets, FID nicely mapped to personal perceived image quality, but as the metric is computed based on summary statistics comparison between features extracted from an inception v3 model, it might be uninformative for other types of datasets, especially the more the metric value tends to zero.I still suggest to avoid visual comparison across short training periods, as our subjective approximate judgment can often be biased by the specific examples we look at or even by the noise contribution alone, and the network might be improving despite no discernible change in the fake samples.
On the other hand, via visual inspection, one can easily spot issues like mode collapse, as illustrated in the following fake samples generated for the Mucha dataset.
This an example of partial mode collapse, where the generator is reduced to generate just a few examples, losing in terms of variety. A prominent cause for this might have been the small size of the dataset, plus common duplicates entries.
See instead the following sample for Franzetta art, especially in relation to the training curves. While some entries might give the impression of mode collapse, it is instead more likely just the fault of way too many duplicates in the training dataset for that specific example.
A decent dataset should generally provide a nice training curve following a Pareto distribution. The official README specifies training times for different configurations, but as said you will obtain most of the improvements during the very first part of your training; the tricky and tedious part is trying to reach the finer quality details.
Plateau is a common situation, where the model is subject to little or no improvement. Another possible but less common behavior is divergence or explosion, where after an initial step improvement the model starts deteriorating and exponentially degenerated to pure noise results.
The action to take for such situations is to retrain from a previous checkpoint (as explained in the previous fine-tuning section), reasonably tweaking the hyperparameters according to the recognized problem. As often mentioned, this is an art by itself, and strongly based on intuition, experience and luck.Many possibilities are available, depending on goals, model status, etc.. Here some of the most common suggestions, accompanied by some empirical notes.
dataset, as usual if your dataset has clear deficits in terms of size or noise, put some effort into fixing those instead of risking to waste time on model tuning
learning-rate (lr), here you can decrease lr for potential more stable (but slow) training progress. You can also tweak the lr of discriminator and generator individually, based on the unbalance that might be present in your architecture between the two components. In this regard, Shawn Presser pointed out on Twitter that v2 has actually a bug, for which the generator lr value is incorrectly used as the discriminator lr. Be sure to patch this part before experimenting with G/D lr balance.I tried to increase lr, just for the laughs, and it always ended up with quick and abrupt divergence.
batch size, increasing batch size can also provide a more stable training or breakthrough from a local minimum. This is based on the idea that as more predictions are used to compute the gradient, the more accurate the weights update and training direction will be. Multiple experiments I run with bigger minibatch forced a good point percentage increase in FID50k, but some of these became soon unstable and exploded. Also, be aware of memory constraints for your setup, v2 is more expensive, and training on 1024 resolution requires a GPU with at least 16 GB of memory.It is worth pointing out that StyleGAN has two different parameters for batch size, minibatch_size_base and minibatch_gpu_base . The former is the actual maximum minibatch size to use, while the latter is the size as processed at a time by a single GPU. This implies that the two parameters need to be adjusted according to the number of GPUs of your setup, as gradients are accumulated until the base minibatch size is reached.
Regarding these aspects, I want to suggest again Gwern StyleGAN v1 post, as it provides in-depth details and references, and three additional excellent papers if you want to explore GAN-training further: Stabilizing Generative Adversarial Network Training: A Survey (2019), Which Training Methods for GANs do actually Converge? (2018), Improved Techniques for Training GANs (2016).
Some behaviors still puzzle me, like what causes this change in loss scores after I restart training from a previous checkpoint (same config).
Above some examples of latent space exploration for my dresses and footwear models. The idea is just to generate N sample vectors (drawn from a Gaussian distribution) and transition between them sequentially using whatever preferred transition function. In my case here this function is just a linear interpolation done on a fixed frames number (equivalent to morphing).
Notice that we are relying on the very initial latent vector z. This means that we are using the StyleGAN mapping network to first generate the latent vector w, and then using w to synthesize a new image. For this reason, we can rely on the truncation trick, and discard areas of the latent space poorly represented. We want to specify how much the generated intermediate vector w has to stay close to the average (computed based on random inputs to the mapping network). ψ (psi) value scales the deviation of w from the average, and as such can be tweaked for quality/variety trade-offs. ψ=1 is equivalent to no truncation (original w), while values towards 0 gets us closer to the average, with quality improvement but a reduction in terms of variety.
We often want to be able to obtain the code/latent-vector/embedding of real images with regards to a target model, in other words: what is the input value I should feed to my model to generate the best approximation of my image.
In general, there are two methods for this:
pass image through the encoder component of the network
optimize latent (using gradient descent)
The former provides a fast solution but has problems generalizing outside of the training dataset, and unfortunately for us, it doesn’t come out of the box with a vanilla StyleGAN. The architecture simply doesn’t learn an explicit encoding function.
We are left with latent optimization option using perceptual loss. We extract high-level features (e.g. from a pre-trained model like VGG) for the reference and generated images, compute the distance between them and optimize on the latent representation (our target code). The initialization of this target code is a very important aspect for efficiency and effectiveness. The easiest way is simple random initialization, but a lot can be done to improve on this, for example by learning an explicit encoding function from images to latent. The idea is to randomly generate a set of N examples and store both the resulting image and the code that generated it. We can then train a model (e.g. ResNet) on this data, and use it to initialize our latent before the actual StyleGAN encoding process. See this rich discussion regarding improved initialization.
Encoder for v1 and Encoder for v2 provide code and step-by-step guide for this operation. I also suggest the following two papers: Image2StyleGAN and Image2StyleGAN++, which give a good overview of encoding images for Stylegan, with considerations about initialization options and latent space quality, plus an analysis of image editing operations like morphing and style mixing.
StyleGAN uses a mapping network (eight fully connected layers) to convert the input noise (z) to an intermediate latent vector (w). Both are of size 512, but the intermediate vector is replicated for each style layer. For a network trained on 1024 size images, this intermediate vector will then be of shape (512, 18), for 512 size it will be (512, 16).
The encoding process is generally done on this intermediate vector, and as such one can decide whether to optimize for w(1) (meaning only one 512 layers, which is then tiled as necessary to each style layer) or the whole w(N). The official projector operated the former, while adaptations often rely on optimizing all w entries individually, for visual fidelity. Regarding this topic see also this Twitter thread.
Even more noticeable and goliardic when projecting reference not proper of the model training distribution, like in the following example projecting a dress image for the FFHQ model.
In general, one can always notice that for high resolutions the projector seems to fail to match fine details of the reference picture, but this is most likely a consequence of using a resolution of 256x256 for the perceptual loss, as demonstrated in this thread.
The StyleGAN improvements on latent space disentanglement allow to explore single attributes of the dataset in a pleasing, orthogonal way (meaning without affecting other attributes).While discriminative models learn boundaries to separate target attributes (e.g. male/female, smile/not-smile, cat/dog) what we are interested here is to cross those boundaries, moving perpendicularly to them. For example, if I start from a sad face I can slowly but steadily move to a smiling version of the same face.
This should already provide a hint on how to learn new directions. We first collect multiple samples (image + latent) from our model and manually classify the images for our target attribute (e.g. smiling VS not smiling), trying to guarantee proper class representation balance. We then train a model to classify or regress on our latents and manual labels. At this point we can use the learned functions of these support models as transition directions.
Robert Luxemburg shared the learned direction for the official FFHQ model.
StyleGAN2 Distillation for Feed-forward Image Manipulation is a very recent paper exploring direction manipulation via a “student” image-to-image network trained on unpaired dataset generated via StyleGAN. The paper aims to overcome the encoding performance-bottleneck and learn a transformation function than can be efficiently applied to real-world images.
A lot of my experiments have initially been motivated by evaluating how good is the latent space learned by a StyleGAN model (representation learning), and how performant the obtained embedding can be for downstream tasks (e.g. image classification via linear models). From one side I will keep working on this kind of evaluation, trying to cover also other generative models types, like autoregressive and flow-based models.
I’m also interested in pursuing the exploration of such models for the pure image-synthesis capabilities, and the increasing potential for effective semantic mixing and editing, especially in relation to my passion for drawing, digital painting and animation. Automated lineart colorization, animating paintings, frames interpolation are some awesome free utils already out there, but there is so much more that can be done for assisted drawing, especially from a more semantic point of view. There is also plenty of room for practical improvements: generalization capabilities, speeding up inference time, training optimization and transfer learning.
Down the line, I also want to go beyond the pure 2-dimensional canvas, and learn more about the amazing things already achieved in the 3D graphics realm. Denoising is something I now frequently rely on, differentiable rendering just blew my mind and to close the loop, back again to GAN for continuous 3D shape generation.
Quoting research scientist Károly Zsolnai-Fehér
“What a time to be alive”
You can check out my code on Github, follow more of my experiments and explanations on Twitter and see my graphics results on Instagram.
Disclaimer: at the moment I cannot share datasets and trained models. Apologies.