<jupyter_start><jupyter_text>What is this?This is a demo to walk people through the results of training StyleGAN2 on various small datasets and under various configurations.You may have seen or be interested in the results generated from the highest resolution pretrained StyleGAN2 models at: https://thispersondoesnotexist.com/Explanations of the effects of certain critical techniques are provided. Most of these models were trained (at the time of writing this) for only about 6 hours on a single P100 GPU via Colab (typically with 4 models training on separate P100s simultaneously, for absolutely no cost). At the end, some results are shown from training for ~20 hours. Typically, models of this nature would be trained up to 100x as long. These will be trained further and this demo will be updated frequently in the near future. However, the results are already decent enough for the purposes of demonstration and comparison.All of the models were trained using a notebook I wrote to allow for easy training using the free resources provided by Colab and Google Drive, linked here:https://colab.research.google.com/drive/1prEbP9AgZnxGCXtZkP-pgqRJoHcHJPou?usp=sharingUsing the PyTorch implementation of StyleGAN2 available at:https://github.com/lucidrains/stylegan2-pytorchThe GitHub repo I made for this project is available at:https://github.com/96jonesa/StyleGan2-Colab-DemoThe public directory on my Google Drive containing all sample images used in this notebook:https://drive.google.com/drive/folders/1gpZKmuvOnsuRmCo3MEcpST_WC1Laaz3W?usp=sharing How to use:Login to Google (Drive).Image displays are full size 8x8 grids of generated imagesYou can either run all the cells then scroll through with Cmd/Ctrl+F9 or 'Runtime > Run all'Or you can step through cell-by-cell by running each cell individually with Cmd/Ctrl+Enter (runs current cell) or Shift+Enter (runs current cell and moves focus to next cell)  Some cool results to get you interestedAll of these images were generated by the trained models - none of these are real photos. Continue through the demo to see how various techniques work together to achieve this quality of generated image (starting from low quality results and progressively adding techniques and improving quality).  Descriptions of datasets (citations at end of notebook):metfaces: "image dataset of human faces extracted from works of art" -NVlabsceleba: "large-scale face attributes dataset with more than 200K celebrity images" -Lie et al.afhq_dog: "5,000 high-quality images [of dogs] at 512×512 resolution" -Choi et al.cifar10_horse: 5000 32x32 color images of horses -Krizhevsky  How does StyleGAN2 work (somewhat technically)?GANs typically have a generator network and a discriminator network. The discriminator network is a glorified classification network - it seeks to determine if a given image came from the underlying data distribution under consideration (i.e. does it look like real data from the dataset being used for training). The generator network seeks to produce images that the discriminator will classify as having come from the underlying data distribution, however it does not get to look at the images in the dataset - it must improve based only on how the discriminator reacts to samples it produces. In order to generate diverse images, generator networks are given random noise as input. Typically during training, the discriminator network is fed some real data, then some generated data, then optimized via stochastic gradient descent. The generator network is optimized after getting the results back from the discriminator network, with fancy scheduling regarding how often each network is optimized (to induce stabilization) depending on the GAN architecture.In StyleGAN2, the generator network is composed of a mapping network and a synthesis network. The mapping network is composed of a sequence of fully-connected layers which serve to transform the random noise input (latent code) into a vector in an intermediate latent space in which the factors of variation are more linear than in the original latent space (i.e. it maps the latent code into a space better topologically structured for the problem). Between the mapping and synthesis networks, affine transformations are learned to produce styles from the intermediate latent vector. These styles are fed as inputs to layers in modules of the synthesis network corresponding to different resolutions (from 4x4 up to the desired resolution). There are several convolutional (each along with standard deviation modulation and normalization) style blocks per resolution module. Per-style block scaling factors are also learned and applied to the noise before being input to intermediate layers of the corresponding style block. After a few style blocks, upsampling occurs to bump the network up to the next resolution (4x4 to 8x8 to 16x16 and so on) until the desired resolution is attained.StyleGAN2 also uses an exponential moving average of the generator (synthesis and mapping) network parameters to compensate for the oscillatory tendencies of GANs, as well as mixed regularities (occasionally using two latent space codes instead of one and simply switching from the intermediate latent vector produced by one to the intermediate latent vector produced by the other at a randomly selected point in the synthesis network, essentially mixing the styles which would be generated by the two). This difference in generated samples due to the addition of these techniques is demonstrated below.There has also been a very recent development in data augmentation that applies to GANs in general (differentiable augmentation), reducing the amount of data required to obtain high-quality results by (in some cases) orders of magnitude. The effect of applying this technique is demonstrated below.  Download the generated images:<jupyter_code># Utilities for downloading publicly shared Google Drive files (from my Google Drive).

import requests

def download_file_from_google_drive(id, destination):
    URL = 'https://docs.google.com/uc?export=download'

    session = requests.Session()

    response = session.get(URL, params = { 'id' : id }, stream = True)
    token = get_confirm_token(response)

    if token:
        params = { 'id' : id, 'confirm' : token }
        response = session.get(URL, params = params, stream = True)

    save_response_content(response, destination)    

def get_confirm_token(response):
    for key, value in response.cookies.items():
        if key.startswith('download_warning'):
            return value

    return None

def save_response_content(response, destination):
    CHUNK_SIZE = 32768

    with open(destination, 'wb') as f:
        for chunk in response.iter_content(CHUNK_SIZE):
            if chunk: # filter out keep-alive new chunks
                f.write(chunk)
import zipfile
from IPython.display import Image, display

file_id = '1nm-yOTaibATgFlM_5rOsA4TB8wxRVAHg'
destination = 'demo_samples.zip'
download_file_from_google_drive(file_id, destination)
zip_ref = zipfile.ZipFile('demo_samples.zip', 'r')
zip_ref.extractall('demo_samples')
zip_ref.close()<jupyter_output><empty_output><jupyter_text>First things firstFirst, we will look at sample images generated by the pretrained models using no differentiable augmentation and no attention. We expect the results to lack diversity (due to the lack of differentiable augmentation) and be of relatively poor quality (due to the limited training time, lack of differentiable augmentation, and lack of use of certain training techniques described further below). In order, these 8x8 image displays are from training on the following datasets:1. cifar10_horse  (32x32 resolution, 5000 images, 97000 iterations @ 4.52 it/s on P100 GPU)2. afhq_dog       (128x128 resolution, 5000 images, 38000 iterations @ 1.78 it/s on P100 GPU)3. metfaces       (128x128 resolution, 1336 images, 27000 iterations @ 1.27 it/s on P100 GPU)4. celeba         (128x128 resolution, 202599 images, 36000 iterations @ 1.64 it/s on P100 GPU) cifar10_horse, no differentiable augmentation, no attention, standard:<jupyter_code>display(Image('/content/demo_samples/StyleGan2_small_set_demo_samples/cifar10_horse_aug_00_attn_none_97.jpg'))<jupyter_output><empty_output><jupyter_text>afhq_dog, no differentiable augmentation, no attention, standard:<jupyter_code>display(Image('/content/demo_samples/StyleGan2_small_set_demo_samples/afhq_dog_aug_00_attn_none_38.jpg'))<jupyter_output><empty_output><jupyter_text>metfaces, no differentiable augmentation, no attention, standard:<jupyter_code>display(Image('/content/demo_samples/StyleGan2_small_set_demo_samples/metfaces_aug_00_attn_none_27.jpg'))<jupyter_output><empty_output><jupyter_text>celeba, no differentiable augmentation, no attention, standard:<jupyter_code>display(Image('/content/demo_samples/StyleGan2_small_set_demo_samples/celeba_aug_00_attn_none_36.jpg'))<jupyter_output><empty_output><jupyter_text>EMA parametersNotice the poor quality? The model also generates samples from a version of the model which use an exponential moving average of the generator (synthesis and mapping) network parameters to compensate for the oscillatory tendencies of GANs. This improves stability and quality significantly. The result is shown below for the same models as above: cifar10_horse, no differentiable augmentation, no attention, EMA parameters:<jupyter_code>display(Image('/content/demo_samples/StyleGan2_small_set_demo_samples/cifar10_horse_aug_00_attn_none_97-ema.jpg'))<jupyter_output><empty_output><jupyter_text>afhq_dog, no differentiable augmentation, no attention, EMA parameters:<jupyter_code>display(Image('/content/demo_samples/StyleGan2_small_set_demo_samples/afhq_dog_aug_00_attn_none_38-ema.jpg'))<jupyter_output><empty_output><jupyter_text>metfaces, no differentiable augmentation, no attention, EMA parameters:<jupyter_code>display(Image('/content/demo_samples/StyleGan2_small_set_demo_samples/metfaces_aug_00_attn_none_27-ema.jpg'))<jupyter_output><empty_output><jupyter_text>celeba, no differentiable augmentation, no attention, EMA parameters:<jupyter_code>display(Image('/content/demo_samples/StyleGan2_small_set_demo_samples/celeba_aug_00_attn_none_36-ema.jpg'))<jupyter_output><empty_output><jupyter_text>Mixed regularitiesThe model also generates samples using EMA generator (synthesis and mapping) parameters as well as mixed regularities (using two latent space codes instead of one and simply switching from the intermediate latent vector produced by one to the intermediate latent vector produced by the other at a randomly selected point in the synthesis network, essentially mixing the styles which would be generated by the two), which causes a decorrelation of neighboring styles and thus allows for more fine-tuned diversity (however, I have found that when under-trained this leads to lower perceptual quality, presumably because the model no longer over-focuses on each distinct style). The images are presented as interpolations between styles both from left to right and from top to bottom: cifar10_horse, no differentiable augmentation, no attention, mixed regularities:<jupyter_code>display(Image('/content/demo_samples/StyleGan2_small_set_demo_samples/cifar10_horse_aug_00_attn_none_97-mr.jpg'))<jupyter_output><empty_output><jupyter_text>afhq_dog, no differentiable augmentation, no attention, mixed regularities:<jupyter_code>display(Image('/content/demo_samples/StyleGan2_small_set_demo_samples/afhq_dog_aug_00_attn_none_38-mr.jpg'))<jupyter_output><empty_output><jupyter_text>metfaces, no differentiable augmentation, no attention, mixed regularities:<jupyter_code>display(Image('/content/demo_samples/StyleGan2_small_set_demo_samples/metfaces_aug_00_attn_none_27-mr.jpg'))<jupyter_output><empty_output><jupyter_text>celeba, no differentiable augmentation, no attention, mixed regularities:<jupyter_code>display(Image('/content/demo_samples/StyleGan2_small_set_demo_samples/celeba_aug_00_attn_none_36-mr.jpg'))<jupyter_output><empty_output><jupyter_text>Early stoppingBy now you have likely noticed the extremely low quality of the afhq_dog model. This model has drifted too far from proper convergence due to repeated failures to convince the discriminator with samples that were moving in the right direction, so it has given in to pressures away from the correct direction and is now repeatedly outputting garbage. In a soon-to-come update to this demo, the use of lower learning rates will be explored as a means of avoiding this issue. Here are the EMA generator (synthesis and mapping) parameters and mixed regularities outputs from an earlier training checkpoint of the same model (22000 iterations instead of 38000 iterations): afhq_dog early stop, no differentiable augmentation, no attention, EMA parameters:<jupyter_code>display(Image('/content/demo_samples/StyleGan2_small_set_demo_samples/afhq_dog_aug_00_attn_none_22-ema.jpg'))<jupyter_output><empty_output><jupyter_text>afhq_dog early stop, no differentiable augmentation, no attention, mixed regularities:<jupyter_code>display(Image('/content/demo_samples/StyleGan2_small_set_demo_samples/afhq_dog_aug_00_attn_none_22-mr.jpg'))<jupyter_output><empty_output><jupyter_text>Differentiable augmentationBy now you have likely noticed the extreme lack of diversity in the afhq_dog and metfaces models, as well as a more subtle lack of diversity in the cifar10_horses and celeba models. This is due to the small amount of data available in the former two datasets, and a reasonable yet still lacking amount of data in the latter two datasets. A recent innovation in data augmentation which uses differentiable augmentations of the data has led to models obtaining high quality results with up to 70x less data. As long as the probability of applying these augmentations to each training sample is kept low enough, the generated images will not contain artifacts of the augmentations (e.g. warping, cropped removals, etc.). The following show the results of training the above models while applying this augmentation to the discriminator input with probability 0.2 at each iteration (the EMA generator (synthesis and mapping) parameters results are shown, without mixed regularities). In order, these were trained on:1. cifar10_horse  (32x32 resolution, 5000 images, 97000 iterations @ 4.51 it/s on P100 GPU)2. afhq_dog       (128x128 resolution, 5000 images, 38000 iterations @ 1.56 it/s on P100 GPU)3. metfaces       (128x128 resolution, 1336 images, 27000 iterations @ 1.05 it/s on P100 GPU)4. celeba         (128x128 resolution, 202599 images, 36000 iterations @ 1.64 it/s on P100 GPU) cifar10_horse, 0.2 differentiable augmentation, no attention, EMA parameters:<jupyter_code>display(Image('/content/demo_samples/StyleGan2_small_set_demo_samples/cifar10_horse_aug_02_attn_none_97-ema.jpg'))<jupyter_output><empty_output><jupyter_text>afhq_dog, 0.2 differentiable augmentation, no attention, EMA parameters:<jupyter_code>display(Image('/content/demo_samples/StyleGan2_small_set_demo_samples/afhq_dog_aug_02_attn_none_38-ema.jpg'))<jupyter_output><empty_output><jupyter_text>metfaces, 0.2 differentiable augmentation, no attention, EMA parameters:<jupyter_code>display(Image('/content/demo_samples/StyleGan2_small_set_demo_samples/metfaces_aug_02_attn_none_27-ema.jpg'))<jupyter_output><empty_output><jupyter_text>celeba, 0.2 differentiable augmentation, no attention, EMA parameters:<jupyter_code>display(Image('/content/demo_samples/StyleGan2_small_set_demo_samples/celeba_aug_02_attn_none_36-ema.jpg'))<jupyter_output><empty_output><jupyter_text>Cool, but let's see that with mixed regularitiesThe same augmented data models produced the following sample images using mixed regularities: cifar10_horse, 0.2 differentiable augmentation, no attention, mixed regularities:<jupyter_code>display(Image('/content/demo_samples/StyleGan2_small_set_demo_samples/cifar10_horse_aug_02_attn_none_97-mr.jpg'))<jupyter_output><empty_output><jupyter_text>afhq_dog, 0.2 differentiable augmentation, no attention, mixed regularities:<jupyter_code>display(Image('/content/demo_samples/StyleGan2_small_set_demo_samples/afhq_dog_aug_02_attn_none_38-mr.jpg'))<jupyter_output><empty_output><jupyter_text>metfaces, 0.2 differentiable augmentation, no attention, mixed regularities:<jupyter_code>display(Image('/content/demo_samples/StyleGan2_small_set_demo_samples/metfaces_aug_02_attn_none_27-mr.jpg'))<jupyter_output><empty_output><jupyter_text>celeba, 0.2 differentiable augmentation, no attention, mixed regularities:<jupyter_code>display(Image('/content/demo_samples/StyleGan2_small_set_demo_samples/celeba_aug_02_attn_none_36-mr.jpg'))<jupyter_output><empty_output><jupyter_text>Longer trainingAll of the results thus far have been from models trained for only about 6 hours each - in order to obtain better results we must train longer. The following show the results of training the above models while applying differentiable augmentation to the discriminator input with probability 0.2 at each iteration (same as above), but this time trained more than three times as long. The EMA results are shown, and are of significantly better quality than the results from the models in their less trained states. In order, these were trained on:1. cifar10_horse (32x32 resolution, 5000 images, 314000 iterations @ 4.51 it/s on P100 GPU)2. afhq_dog (128x128 resolution, 5000 images, 119000 iterations @ 1.56 it/s on P100 GPU)3. metfaces (128x128 resolution, 1336 images, 88000 iterations @ 1.05 it/s on P100 GPU)4. celeba (128x128 resolution, 202599 images, 120000 iterations @ 1.64 it/s on P100 GPU) cifar10_horse, 0.2 differentiable augmentation, no attention, EMA parameters:<jupyter_code>display(Image('/content/demo_samples/StyleGan2_small_set_demo_samples/cifar10_horse_aug_02_attn_none_314-ema.jpg'))<jupyter_output><empty_output><jupyter_text>afhq_dog, 0.2 differentiable augmentation, no attention, EMA parameters:<jupyter_code>display(Image('/content/demo_samples/StyleGan2_small_set_demo_samples/afhq_dog_aug_02_attn_none_119-ema.jpg'))<jupyter_output><empty_output><jupyter_text>metfaces, 0.2 differentiable augmentation, no attention, EMA parameters:<jupyter_code>display(Image('/content/demo_samples/StyleGan2_small_set_demo_samples/metfaces_aug_02_attn_none_88-ema.jpg'))<jupyter_output><empty_output><jupyter_text>celeba, 0.2 differentiable augmentation, no attention, EMA parameters:<jupyter_code>display(Image('/content/demo_samples/StyleGan2_small_set_demo_samples/celeba_aug_02_attn_none_120-ema.jpg'))<jupyter_output><empty_output><jupyter_text>And the mixed regularities results... cifar10_horse, 0.2 differentiable augmentation, no attention, mixed regularities:<jupyter_code>display(Image('/content/demo_samples/StyleGan2_small_set_demo_samples/cifar10_horse_aug_02_attn_none_314-mr.jpg'))<jupyter_output><empty_output><jupyter_text>afhq_dog, 0.2 differentiable augmentation, no attention, mixed regularities:<jupyter_code>display(Image('/content/demo_samples/StyleGan2_small_set_demo_samples/afhq_dog_aug_02_attn_none_119-mr.jpg'))<jupyter_output><empty_output><jupyter_text>metfaces, 0.2 differentiable augmentation, no attention, mixed regularities:<jupyter_code>display(Image('/content/demo_samples/StyleGan2_small_set_demo_samples/metfaces_aug_02_attn_none_88-mr.jpg'))<jupyter_output><empty_output><jupyter_text>celeba, 0.2 differentiable augmentation, no attention, mixed regularities:<jupyter_code>display(Image('/content/demo_samples/StyleGan2_small_set_demo_samples/celeba__aug_02_attn_none_120-mr.jpg'))<jupyter_output><empty_output><jupyter_text>Longer training without differentiable augmentationThe following shows the results of continued training on the models without differentiable augmentation. The EMA results are shown, and although they have improved on the results from the less trained models without augmentation, they are much worse than the results from further training with augmentation. The models (except for celeba, which has a much larger dataset) have collapsed into single representations of poor quality, and the afhq_dog model failed to stabilize after several attempts at training. In the case of cifar10_horse, there appears to be a case of overtraining due to the smaller image size - expect an update soon demonstrating and explaining the deterioration of the quality from its peak. In order, these were trained on:1. cifar10_horse (32x32 resolution, 5000 images, 273000 iterations @ 4.51 it/s on P100 GPU)2. afhq_dog (128x128 resolution, 5000 images, 76000 iterations @ 1.56 it/s on P100 GPU)3. metfaces (128x128 resolution, 1336 images, 81000 iterations @ 1.05 it/s on P100 GPU)4. celeba (128x128 resolution, 202599 images, 112000 iterations @ 1.64 it/s on P100 GPU) cifar10_horse, no differentiable augmentation, no attention, EMA parameters:<jupyter_code>display(Image('/content/demo_samples/StyleGan2_small_set_demo_samples/cifar10_horse_aug_00_attn_none_273-ema.jpg'))<jupyter_output><empty_output><jupyter_text>afhq_dog, no differentiable augmentation, no attention, EMA parameters:<jupyter_code>display(Image('/content/demo_samples/StyleGan2_small_set_demo_samples/afhq_dog_aug_00_attn_none_76-ema.jpg'))<jupyter_output><empty_output><jupyter_text>metfaces, no differentiable augmentation, no attention, EMA parameters:<jupyter_code>display(Image('/content/demo_samples/StyleGan2_small_set_demo_samples/metfaces_aug_00_attn_none_81-ema.jpg'))<jupyter_output><empty_output><jupyter_text>celeba, no differentiable augmentation, no attention, EMA parameters:<jupyter_code>display(Image('/content/demo_samples/StyleGan2_small_set_demo_samples/celeba_aug_00_attn_none_112-ema.jpg'))<jupyter_output><empty_output><jupyter_text>And the mixed regularities results without augmentation... cifar10_horse, no differentiable augmentation, no attention, mixed regularities:<jupyter_code>display(Image('/content/demo_samples/StyleGan2_small_set_demo_samples/cifar10_horse_aug_00_attn_none_273-mr.jpg'))<jupyter_output><empty_output><jupyter_text>afhq_dog, no differentiable augmentation, no attention, mixed regularities:<jupyter_code>display(Image('/content/demo_samples/StyleGan2_small_set_demo_samples/afhq_dog_aug_00_attn_none_76-mr.jpg'))<jupyter_output><empty_output><jupyter_text>metfaces, no differentiable augmentation, no attention, mixed regularities:<jupyter_code>display(Image('/content/demo_samples/StyleGan2_small_set_demo_samples/metfaces_aug_00_attn_none_81-mr.jpg'))<jupyter_output><empty_output><jupyter_text>celeba, no differentiable augmentation, no attention, mixed regularities:<jupyter_code>display(Image('/content/demo_samples/StyleGan2_small_set_demo_samples/celeba_aug_00_attn_none_112-mr.jpg'))<jupyter_output><empty_output><jupyter_text>Higher resolution (256x256)An obvious bound on the quality of images is the image resolution - so far we have only generated 128x128 images (and 32x32 in the case of cifar10_horse). However, the afhq_dog and metfaces datasets are 1024x1024, so we can train models on them to generate images up to that resolution. Here we present the results from training at 256x256 resolution with differentiable augmentation (first the EMA results) - notice how the images appear to be warped in some cases (especially in the metfaces model), a result of the differential augmentation bleeding through to the generator (this can and will be fixed by training for a longer time and/or decreasing the augmentation probability):1. afhq_dog (256x256 resolution, 5000 images, 72000 iterations @ 1.10 s/it on P100 GPU)2. metfaces (256x256 resolution, 1336 images, 68000 iterations @ 1.10 s/it on P100 GPU) afhq_dog 256x256, 0.2 differentiable augmentation, no attention, EMA parameters:<jupyter_code>display(Image('/content/demo_samples/StyleGan2_small_set_demo_samples/afhq_dog_size_256_aug_02_attn_none_72-ema.jpg'))<jupyter_output><empty_output><jupyter_text>metfaces 256x256, 0.2 differentiable augmentation, no attention, EMA parameters:<jupyter_code>display(Image('/content/demo_samples/StyleGan2_small_set_demo_samples/metfaces_size_256_aug_02_attn_none_68-ema.jpg'))<jupyter_output><empty_output><jupyter_text>And now the mixed regularities results afhq_dog 256x256, 0.2 differentiable augmentation, no attention, mixed regularities:<jupyter_code>display(Image('/content/demo_samples/StyleGan2_small_set_demo_samples/afhq_dog_size_256_aug_02_attn_none_72-mr.jpg'))<jupyter_output><empty_output><jupyter_text>metfaces 256x256, 0.2 differentiable augmentation, no attention, mixed regularities:<jupyter_code>display(Image('/content/demo_samples/StyleGan2_small_set_demo_samples/metfaces_size_256_aug_02_attn_none_68-mr.jpg'))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.06482942818987496, 'fraction_numerical': 0.03520699283719801, 'mean_word_length': 7.123602892833662, 'pattern_counts': {'":': 0, '<': 173, '<?xml version=': 0, '>': 174, 'https://': 6, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '43', 'dir': 'jupyter-structured-clean-dedup', 'id': '170234', 'license': 'MIT', 'n_tokens_mistral': 7770, 'n_tokens_neox': 6953, 'n_words': 2353, 'path': 'Ipynb/output_small_set_demo.ipynb', 'repo_name': 'punyaslokdutta/StyleGAN_FacialFeatureEditor', 'size': '24711'}