Each example is designed to be self-contained and easily forkable, while reproducing relevant results in different areas of machine learning. These examples adhere to a shared style and functionality outlined in #231. All examples are under the folder flax/examples/. Some of the examples below have a link [Interactive] that lets you run them directly in Colab.
Image classification
- MNIST [Interactive] : Convolutional neural network for MNIST classification (featuring simple code).
- ImageNet : Resnet-50 on imagenet with weight decay (featuring multi host SPMD, custom preprocessing, checkpointing, dynamic scaling, mixed precision).
Reinforcement Learning
- Proximal Policy Optimization : Learning to play Atari games (featuring single host SPMD, RL setup).
Natural language processing
- Sequence to sequence for number addition (featuring simple code, LSTM state handling, on the fly data generation).
- Transformer model on WMT : Translating English/German (featuring multihost SPMD, dynamic bucketing, attention cache, packed sequences, recipe for TPU training on GCP).
Generative models
- Variational auto-encoder : Trained on binarized MNIST (featuring simple code, vmap).
- PixelCNN++ : Trained on cifar10 (featuring single host SPMD, checkpointing, Polyak decay).
In addition to the curated list of official Flax examples, there is a growing community of people using Flax to build new types of machine learning models. We are happy to showcase any example built by the community here! If you want to submit your own example, we suggest that you start by forking one of the official Flax example, and start from there.
| Link | Author | Task type | Reference |
|---|---|---|---|
| JAX-RL | @henry-prior | Reinforcement learning | N/A |
| DCGAN Colab | @bkkaggle | Image Synthesis | https://arxiv.org/abs/1511.06434 |
Using the Deprecated flax.nn API ~~~~~~~~~~~~~~~~~~~~
The following examples were created using the old pre-Linen API. You can still look at them for inspiration. Or maybe ask the authors if they would accept a pull request to the new API if you want to earn some Flax karma?
| Link | Author | Task type | Reference |
|---|---|---|---|
| Gaussian Processes regression | @danieljtait | Regression | N/A |
| DQN | @joaogui1 | Reinforcement learning | N/A |
| Various CIFAR SOTA Models | @PForet | Image Classification | N/A |
Looking for "FOO" implemented in Flax? We use GitHub issues to keep track of which models people are most interested in seeing re-implemented in Flax. If you can't find what you're looking for in the list "example requested", file an issue with this template. If the model you are looking for has already been requested by others, upvote the issue to help us see which ones are the most requested.
Looking to implement something in Flax? Consider looking at the list "example requested" and go ahead and build it!