Guidelines for Examples

Following some discussion I wrote down a set of guidelines for Flax examples inside the google/flax repository. It will take some work to update the current set of examples but I believe that the examples will be even better afterwards and easier to maintain.

Guidelines

• Each example should have at least 2 owners . Creating examples cost time and we want to recognize this. We also would like to have a point of contact in case the example needs to be updated to a new API.
• Every example should have a README.md which specifies at least:
  • Command for running the example with a link to http://tensorboard.dev for a successful run. Optionally one can also include some output logged on the command line.
  • Link to related references (paper for the model or a great blogpost).
• Keep the example focused.
  • Each example should support a single model on a single dataset.
  • Keep the number of configurable hyper-parameters basic (<10) and use absl.flags. We do not want to overwhelm readers.
• Default hyperparameter configuration should be provided as ConfigDict in configs/default.py. An additional configs/test.py for a CPU friendly unit test case are ok.
• To make it easier to test, maintain and reuse code structure the code as follows:
  • main.py contains the flags and calls a method from train.py to run the training loop. This should be the only file defining flags! This can be almost identical for all examples, please copy from linen_examples/wmt/main.py. If possible stick with the 2 flags (--workdir (not model_dir) and --config).
  • train.py contains classes and methods for training and evaluating the model.
  • train_test.py for test cases of the training code. At a minimum the test should run a single training step but more fine grained unit tests are a bonus. You can use tfds.testing.mock_data to avoid real data from disk.
  • Additional files for more complex architectures or input_pipeline.py.
• Only use public datasets available in TensorFlow Datasets. TensorFlow Datasets standardizes how datasets downloaded and prepared. Synthetic data is also allowed.
• Follow the Google Python Style Guide (no relative imports, use type annotation)
• requirements.txt if running the example requires additional Python packages that are not installed by Flax or it’s dependencies.
• Use the clu package where appropriate. This helps us to de-duplicate common functionality and make the code agnostic to the platform. We would like to see examples that work equally well on Google’s internal systems as on Cloud or locally.

Future Work

Examples should also be covered by regression tests, see #144.

We should also add general instructions for running on examples (on GPU and TPU) using Google Cloud.

[avital]

@mohitreddy1996 has agreed to get the ball rolling on this by porting MNIST and/or ImageNet. We’ll review that and then if a quorum of us like it, we can split up the work for porting the other examples.

That should also help unblock #176

[andsteing]

Closing this issue with #1841