A Jax/Flax codebase for running deterministic, scalable, and well-documented deep learning experiments, with a particular emphasis on neural network initialization, optimization, and tuning experiments.

There is not yet a stable version (nor an official release of this library). All APIs are subject to change.

This is a research project, not an official Google product.

The current development version requires Python 3.6-3.8.

To install the latest development version inside a virtual environment, run

python3 -m venv env-i2w
source env-i2w/bin/activate
pip install --upgrade pip
pip install "git+https://github.com/google/init2winit.git#egg=init2winit"
pip install --upgrade jax jaxlib==0.1.66+cuda111 -f https://storage.googleapis.com/jax-releases/jax_releases.html

where cuda111 corresponds to the installed version of CUDA. For more Jax install information see the Jax README.

An example MNIST experiment can be run with the following command:

python3 main.py \
    --experiment_dir=/tmp/test_mnist \
    --model=fully_connected \
    --dataset=mnist \
    --num_train_steps=10

For local debugging we recommend using the fake dataset:

python3 main.py \
    --experiment_dir=/tmp/test_fake \
    --num_train_steps=10 \
    --dataset=fake \
    --hparam_overrides='{"input_shape": [28, 28, 1], "output_shape": [10]}'

The hparam_overrides accepts a serialized JSON object with hyperparameter names/values to use. See the flags in main.py for more information on possible configurations.

See the dataset_lib and model_lib directories for currently implemented datasets and models.

To cite this repository:

@software{init2winit2021github,
  author = {Justin M. Gilmer and George E. Dahl and Zachary Nado and Priya Kasimbeg and Sourabh Medapati},
  title = {{init2winit}: a JAX codebase for initialization, optimization, and tuning research},
  url = {http://github.com/google/init2winit},
  version = {0.0.2},
  year = {2023},
}

For a list of references to the models, datasets, and techniques implemented in this codebase, see references.md.

Contributors (past and present):

• Ankush Garg
• Behrooz Ghorbani
• Cheolmin Kim
• David Cardoze
• George E. Dahl
• Justin M. Gilmer
• Michal Badura
• Priya Kasimbeg
• Rohan Anil
• Sourabh Medapati
• Sneha Kudugunta
• Varun Godbole
• Zachary Nado
• Vlad Feinberg
• Derrick Xin
• Naman Agarwal
• Daniel Suo
• Bilal Khan
• Jeremy Cohen
• Kacper Krasowiak