learned_optimization: 使用 JAX 进行元学习优化器及其他研究

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

要开始使用，请参阅我们的 文档。

快速入门 Colab 笔记本

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

learned_optimization 教程系列

1. 简介：
2. 创建自定义任务：
3. 截断步数：
4. 梯度估计器：
5. 元训练：
6. 自定义学习型优化器：

从零构建一个学习型优化器

这是一个简单、自包含的学习型优化器示例，不依赖于 learned_optimization 库：

本地安装

我们强烈建议使用 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, 和 Jascha Sohl-Dickstein. "Unbiased gradient estimation in unrolled computation graphs with persistent evolution strategies." 国际机器学习会议（最佳论文奖）。PMLR, 2021。
• Metz, Luke*, C. Daniel Freeman*, Samuel S. Schoenholz, 和 Tal Kachman. "Gradients are Not All You Need." arXiv 预印本 arXiv:2111.05803 (2021)。

开发 / 运行测试

我们将测试文件放置在相关源代码旁边，而不是放在单独的 tests/ 文件夹中。 每个测试都可以直接运行，或者使用 pytest 运行（例如 python3 -m pytest learned_optimization/outer_trainers/）。也可以使用 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},
}

免责声明

learned_optimization 并非 Google 的官方产品。

learned_optimization 快速上手指南

learned_optimization 是一个基于 JAX 的研究代码库，用于训练、设计、评估和应用学习到的优化器（Learned Optimizers），以及更广泛的动态系统元训练。它实现了手工设计和学习到的优化器、元训练/测试任务，以及 ES、PES 和截断时间反向传播等外部训练算法。

环境准备

• 操作系统: Linux 或 macOS (Windows 用户建议使用 WSL2)
• Python 版本: 推荐 Python 3.8+
• 硬件加速: 强烈建议配备 GPU 或 TPU。该库涉及大量矩阵运算和元学习过程，在无加速器环境下运行极慢。
• 前置依赖:
  • git
  • pip
  • virtualenv (推荐用于隔离环境)

国内开发者提示：由于该库深度依赖 JAX 及其底层加速库，安装时若遇到网络问题，可尝试配置国内镜像源（如清华源、阿里源）加速 pip 包下载。但请注意，JAX 的预编译二进制文件（wheels）有时在第三方镜像中更新不及时，若安装失败请切换回官方源。

安装步骤

推荐使用虚拟环境进行安装，以避免依赖冲突。

1. 安装 virtualenv 并克隆仓库：

   pip3 install virtualenv
   git clone https://github.com/google/learned_optimization.git
   cd learned_optimization
   

2. 创建并激活虚拟环境：

   python3 -m venv env
   source env/bin/activate
   

   (Windows 用户使用 env\Scripts\activate)

3. 安装依赖包：

   pip install -e .
   

   (如需加速，可在命令前添加 -i https://pypi.tuna.tsinghua.edu.cn/simple，若遇到 JAX 相关报错请移除该参数重试)

基本使用

方式一：通过 Colab 快速体验（推荐）

对于初次接触或希望快速验证功能的开发者，官方提供了一系列完整的 Jupyter Notebook 教程，支持直接在 Google Colab 中运行（需选择 GPU/TPU 运行时）。

核心教程序列：

1. 入门介绍: Part1_Introduction.ipynb
2. 创建自定义任务: Part2_CustomTasks.ipynb
3. 截断步长 (Truncated Steps): Part3_Truncation_TruncatedStep.ipynb
4. 梯度估计器: Part4_GradientEstimators.ipynb
5. 元训练 (Meta training): Part5_Meta_training_with_GradientLearner.ipynb
6. 自定义学习到的优化器: Part6_custom_learned_optimizers.ipynb

从零构建示例（不依赖库内部高级封装，适合理解原理）：

• no_dependency_learned_optimizer.ipynb

方式二：本地命令行训练示例

安装完成后，您可以直接运行官方提供的示例脚本来训练一个简单的学习到的优化器。

执行命令：

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

说明：

• 该命令会自动使用 tfds 下载所需数据。
• 运行几分钟后，终端将开始打印训练指标。
• 日志保存在 /tmp/logs_folder，可使用 TensorBoard 指向该目录进行可视化：

  tensorboard --logdir=/tmp/logs_folder
  

• 注意：若无 GPU/TPU 加速，此过程将非常缓慢。