Pretrained models in Flax (GPT2, StyleGAN2, ResNets etc)

[matthias-wright]

Hi all,

I created a repository of pretrained models in Flax that can be easily installed via pip.
Github: https://github.com/matthias-wright/flaxmodels

Current models

• GPT2
• StyleGAN2
• ResNet{18, 34, 50, 101, 152}
• VGG{16, 19}

I will also add more models in the future.

Here are some notebooks to play with on Colab
GPT2, StyleGAN2, ResNet, VGG

I hope this will be useful to someone :)

I also would like to thank all the developers and contributors of Jax and Flax for their work.

Cheers!

[avital]

This is super exciting @matthias-wright! What can we do to continue to support your work?

[avital]

(And do you plan to include any training loops as well?)

[matthias-wright]

Hi @avital, thanks for offering your support! I think I have everything I need right now. If you want, maybe you can link to this project from either the Flax repo or docs.

I haven't included any training loops so far because I am still unsure about how to do that exactly.
One way would be to just include a simple trainer for each model (similar to minGPT) that takes in the model, dataloaders, and some training parameters (such as batch size etc).
The upside is that it makes it easy for users to run some basic trainings. The downside is that customized trainings would probably get complicated fast.
An alternative way would be to create a separate repository (or a submodule) that contains training scripts for each model which users can simply clone. Then they can either run the default training or customize the scripts for their specific use cases.

I would love to get your feedback on this! Which option do you think is better and is there another way that I am not seeing?

[andsteing]

This is very cool indeed!

I'd love to add a link from our examples page (I think it would be best discoverable from there).

Maybe you could change a couple of small things in the repo?

• Submitting the Colabs inside the directory where the model codes resides would make it a lot more discoverable. And give them commit history and possibility for other users to contribute etc. Having a simple fine-tuning example in one of those Colabs might be very useful (and later linking to repos using the models from the model garden...)
• Mention how you verify the proper functioning of the models (I saw you compare the output against some golden output, but maybe you could mention that in the README to make it more discoverable).
• Currently there's quite a lot of information in the individual model READMEs. Maybe you could refactor the usage examples into the Colabs, and the API documentation into the actual models so they're viewable through pydoc from Colab & IDEs. (I noticed in e.g. resnet you have the API docs in both the README and as pydocs, but e.g. in gpt2 it's only in the README).
• Documenting the model checkpoints in some more detail would be great (directly copied from the repo mentioned in the README? Preprocessed somehow? If the latter, checking in that code might also be very helpful).

Cheers, Andreas

[matthias-wright]

Hi @andsteing, thanks for the detailed feedback! I have adapted the repo according to your suggestions and I really like the outcome. Adding a link to the examples page sounds great!