Core examples#

Core examples are hosted on the GitHub Flax repository in the examples directory.

Each example is designed to be self-contained and easily forkable, while reproducing relevant results in different areas of machine learning.

As discussed in #231, we decided to go for a standard pattern for all examples including the simplest ones (like MNIST). This makes every example a bit more verbose, but once you know one example, you know the structure of all of them. Having unit tests and integration tests is also very useful when you fork these 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 - Interactive🕹: Resnet-50 on ImageNet with weight decay (featuring multi host SPMD, custom preprocessing, checkpointing, dynamic scaling, mixed precision).

Reinforcement learning#

Natural language processing#

Generative models#

Graph modeling#

Contributing to core Flax examples#

Most of the core Flax examples on GitHub follow a structure that the Flax dev team found works well with Flax projects. The team strives to make these examples easy to explore and fork. In particular (as per GitHub Issue #231):

  • README: contains links to paper, command line, TensorBoard metrics.

  • Focus: an example is about a single model/dataset.

  • Configs: we use ml_collections.ConfigDict stored under configs/.

  • Tests: executable main.py loads train.py which has train_test.py.

  • Data: is read from TensorFlow Datasets.

  • Standalone: every directory is self-contained.

  • Requirements: versions are pinned in requirements.txt.

  • Boilerplate: is reduced by using clu.

  • Interactive: the example can be explored with a Colab.