I have recently tried to find a state-of-the-art conditional image generation model in Tensorflow 2. I found out that the BigGAN deep model from the Large Scale GAN Training for High Fidelity Natural Image Synthesis paper was a perfect cast for the job. In this post, I present a full Tensorflow 2 implementation of the model with pre-trained weights.
Deepmind released the models on TFHub available as a KerasLayer object but they are not fine-tunable using Tensorflow 2 and the model’s code is not released. This was problematic since I wanted the control to fine-tune the first blocks while freezing the last ones.
I found an implementation by HuggingFace that actually recodes the model’s logic in PyTorch and fetch the pre-trained weights from the released models. I didn’t want to switch to PyTorch so I adapted the model’s logic and weights mapping from PyTorch to Tensorflow 2. The code to build and load pre-trained weights is available on our Github TF2 published models repository.
To generate samples, simply import the BigGAN object and sample.
Since the model’s code is entirely release in TF2, you can fine-tune or change it as you please for instance by freezing the lasts blocks while training the first blocks. Enjoy!