Thread by @labmlai

1/4) Simple StyleGAN2 implementation (with notes) in @PyTorch in less than 500 lines of code (including training)

Annotated Code: nn.labml.ai/gan/stylegan/i…

Github: github.com/lab-ml/annotat…

Paper: arxiv.org/abs/1912.04958

@jaakkolehtinen @NVIDIAAI

🧵👇

2/4) StyleGAN2 is based on StyleGAN which is based on ProgressiveGAN (by the same authors). We have briefly discussed these models.

StyleGAN paper: arxiv.org/abs/1812.04948
ProgressiveGAN paper: arxiv.org/abs/1710.10196

👇

3/4) This is a minimalistic Style GAN2 model training code. Only single-GPU training is supported to keep the implementation simple. We managed to shrink it to keep it at 425 lines of code, including the training loop.

We trained it on 128x128 resolution.

👇

3/4) We ran into a lot of subtle bugs while implementing this. Experiments had to run for quite sometime to see if they were working properly by looking at generated images. Monitoring layer outputs/parameters/gradients helped us identify problems early on.

👇

4/4) We included diagrams in the notes to illustrate model structure. It took quite a bit of effort and would love to know if they help you 😊