learned_optimization: 使用JAX进行元学习优化器及更多

learned_optimization 是一个用于训练、设计、评估和应用学习型优化器的研究代码库，并可更广泛地用于动态系统的元训练。它实现了手工设计和学习型优化器、用于元训练和元测试的任务，以及诸如ES、PES和截断反向传播等外部训练算法。

开始使用请查看我们的文档。

快速入门Colab笔记本

我们的文档也可以作为Colab笔记本运行！我们建议在Colab中使用免费加速器（TPU或GPU）运行这些笔记本（转到运行时-> 更改运行时类型）。

learned_optimization 教程系列

简介：<a href="https://colab.research.google.com/github/google/learned_optimization/blob/main/docs/notebooks/Part1_Introduction.ipynb" target="_parent"><img src="https://yellow-cdn.veclightyear.com/835a84d5/afe700ae-d411-40c9-910d-ee43f9a68ce7.svg" alt="在Colab中打开"/></a>
创建自定义任务：<a href="https://colab.research.google.com/github/google/learned_optimization/blob/main/docs/notebooks/Part2_CustomTasks.ipynb" target="_parent"><img src="https://yellow-cdn.veclightyear.com/835a84d5/afe700ae-d411-40c9-910d-ee43f9a68ce7.svg" alt="在Colab中打开"/></a>
截断步骤：<a href="https://colab.research.google.com/github/google/learned_optimization/blob/main/docs/notebooks/Part3_Truncation_TruncatedStep.ipynb" target="_parent"><img src="https://yellow-cdn.veclightyear.com/835a84d5/afe700ae-d411-40c9-910d-ee43f9a68ce7.svg" alt="在Colab中打开"/></a>
梯度估计器：<a href="https://colab.research.google.com/github/google/learned_optimization/blob/main/docs/notebooks/Part4_GradientEstimators.ipynb" target="_parent"><img src="https://yellow-cdn.veclightyear.com/835a84d5/afe700ae-d411-40c9-910d-ee43f9a68ce7.svg" alt="在Colab中打开"/></a>
元训练：<a href="https://colab.research.google.com/github/google/learned_optimization/blob/main/docs/notebooks/Part5_Meta_training_with_GradientLearner.ipynb" target="_parent"><img src="https://yellow-cdn.veclightyear.com/835a84d5/afe700ae-d411-40c9-910d-ee43f9a68ce7.svg" alt="在Colab中打开"/></a>
自定义学习型优化器：<a href="https://colab.research.google.com/github/google/learned_optimization/blob/main/docs/notebooks/Part6_custom_learned_optimizers.ipynb" target="_parent"><img src="https://yellow-cdn.veclightyear.com/835a84d5/afe700ae-d411-40c9-910d-ee43f9a68ce7.svg" alt="在Colab中打开"/></a>

从头开始构建学习型优化器

不依赖于learned_optimization库的简单、自包含的学习型优化器示例： <a href="https://colab.research.google.com/github/google/learned_optimization/blob/main/docs/notebooks/no_dependency_learned_optimizer.ipynb" target="_parent"><img src="https://yellow-cdn.veclightyear.com/835a84d5/afe700ae-d411-40c9-910d-ee43f9a68ce7.svg" alt="在Colab中打开"/></a>

本地安装

我们强烈建议使用virtualenv来使用此包。

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 .

训练学习型优化器示例

要在简单的内部问题上训练学习型优化器，请运行以下命令：

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

这将首先使用tfds下载数据，然后开始运行。几分钟后，您应该会看到打印的数字。

可以将tensorboard指向此目录以可视化结果。请注意，如果没有加速器，运行速度会非常慢。

需要帮助？有问题？

提交GitHub问题！我们会尽最大努力及时回复。

使用learned_optimization的出版物

写了使用learned_optimization的论文或博客文章？将其添加到列表中！

Vicol, Paul, Luke Metz, and Jascha Sohl-Dickstein. "Unbiased gradient estimation in unrolled computation graphs with persistent evolution strategies." 国际机器学习会议（最佳论文奖）。PMLR，2021。
Metz, Luke*, C. Daniel Freeman*, Samuel S. Schoenholz, and Tal Kachman. "Gradients are Not All You Need." arXiv预印本 arXiv:2111.05803 (2021)。

开发 / 运行测试

我们将测试文件放在相关源代码旁边，而不是放在单独的tests/文件夹中。每个测试可以直接运行，也可以使用pytest运行（例如python3 -m pytest learned_optimization/outer_trainers/）。pytest还可以用于运行所有测试，使用python3 -m pytest，但这将花费相当长的时间。

如果发现任何问题，请提交问题，我们会尽快查看！

引用 learned_optimization

引用此存储库：

@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},
}