Training StyleGAN2 in Jax/Flax

[matthias-wright]

Hi all,

Here is an implementation of StyleGAN2 in Jax/Flax that is compatible with the official weights. The training code can be found here.

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

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!

Some results

Style mixing

[8bitmp3]

🔥

@levskaya @andsteing @marcvanzee @jheek @gnecula @avital @mattjj