Pretrained models for Jax/Flax: StyleGAN2, GPT2, VGG, ResNet.

About

The goal of this project is to make current deep learning models more easily available for the awesome Jax/Flax ecosystem.

Models

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

Example Notebooks to play with on Colab

• GPT2
• StyleGAN2
• ResNet
• VGG

Installation

You will need Python 3.7 or later.

1. For GPU usage, follow the Jax installation with CUDA.
2. Then install:

   > pip install --upgrade git+https://github.com/matthias-wright/flaxmodels.git

For CPU-only you can skip step 1.

Documentation

The documentation for the models is on the individual model pages.

Testing

To run the tests, pytest needs to be installed.

> git clone https://github.com/matthias-wright/flaxmodels.git
> cd flaxmodels
> python -m pytest tests/

Acknowledgments

Thank you to the developers of Jax and Flax. The title image is a photograph of a flax flower, kindly made available by Marta Matyszczyk.

License

Each model has an individual license.