/learned_optimization

Primary LanguagePythonApache License 2.0Apache-2.0

learned_optimization: Meta-learning optimizers and more with JAX

Documentation Status License

learned_optimization is a research codebase for training, designing, evaluating, and applying learned optimizers, and for meta-training of dynamical systems more broadly. It implements hand-designed and learned optimizers, tasks to meta-train and meta-test them, and outer-training algorithms such as ES, PES, and truncated backprop through time.

To get started see our documentation.

Quick Start Colab Notebooks

Our documentation can also be run as colab notebooks! We recommend running these notebooks with a free accelerator (TPU or GPU) in colab (go to Runtime -> Change runtime type).

learned_optimization tutorial sequence

  1. Introduction : Open In Colab
  2. Creating custom tasks: Open In Colab
  3. Truncated Steps: Open In Colab
  4. Gradient estimators: Open In Colab
  5. Meta training: Open In Colab
  6. Custom learned optimizers: Open In Colab

Build a learned optimizer from scratch

Simple, self-contained, learned optimizer example that does not depend on the learned_optimization library: Open In Colab

Local Installation

We strongly recommend using virtualenv to work with this package.

pip3 install virtualenv
git clone git@github.com:google/learned_optimization.git
cd learned_optimization
python3 -m venv env
source env/bin/activate
pip install -e .

Train a learned optimizer example

To train a learned optimizer on a simple inner-problem, run the following:

python3 -m learned_optimization.examples.simple_lopt_train --train_log_dir=/tmp/logs_folder --alsologtostderr

This will first use tfds to download data, then start running. After a few minutes you should see numbers printed.

A tensorboard can be pointed at this directory for visualization of results. Note this will run very slowly without an accelerator.

Need help? Have a question?

File a github issue! We will do our best to respond promptly.

Publications which use learned_optimization

Wrote a paper or blog post that uses learned_optimization? Add it to the list!

Development / Running tests

We locate test files next to the related source as opposed to in a separate tests/ folder. Each test can be run directly, or with pytest (e.g. python3 -m pytest learned_optimization/outer_trainers/). Pytest can also be used to run all tests with python3 -m pytest, but this will take quite some time.

If something is broken please file an issue and we will take a look!

Citing learned_optimization

To cite this repository:

@inproceedings{metz2022practical,
  title={Practical tradeoffs between memory, compute, and performance in learned optimizers},
  author={Metz, Luke and Freeman, C Daniel and Harrison, James and Maheswaranathan, Niru and Sohl-Dickstein, Jascha},
  booktitle = {Conference on Lifelong Learning Agents (CoLLAs)},
  year = {2022},
  url = {http://github.com/google/learned_optimization},
}

Disclaimer

learned_optimization is not an official Google product.