Training StyleGAN2 in Flax

[matthias-wright]

Hi all,

a while ago I posted here about some pretrained models that I ported to Flax. I have now added the training code for StyleGAN2: https://github.com/matthias-wright/flaxmodels/tree/main/training/stylegan2

Highlights

• Multi GPU training (data parallelism) with jax.pmap
• Results match the original implementation (see below)
• tf.data.TFRecordDataset is used for efficient data loading
• optax is used for optimization

Lowlights

• I only trained on resolution 512x512 due to limited resources
• Mixed-precision training is not stable yet. I used flax.optim.DynamicScale and also all the tricks from the original implementation (casting to float32 for some operations, using pre-normalization in the modulated conv layer, only using float16 for the higher resolutions, clipping the output of the convolution layers, etc) but it is not properly working yet.

I also added instructions on how to train on your own dataset. If you have any questions feel free to ask them here or on the github page.

Cheers!

P.S. I also added training code for ResNet and VGG but that is less interesting since there are already some great examples here, which I also borrowed from.

Some results

Style mixing