Optax

简介

Optax 是一个用于 JAX 的梯度处理与优化库。

Optax 旨在通过提供可灵活组合的构建模块来促进研究工作。

我们的目标是：

• 提供简单、经过充分测试且高效的底层组件实现。
• 通过允许用户轻松地将低层次的组件组合成自定义优化器（或其他梯度处理组件），提高研究效率。
• 通过使任何人都能轻松贡献，加速新想法的采用。

我们倾向于专注于小型、可组合的构建模块，这些模块可以有效地组合成自定义解决方案。其他开发者也可以在这些基础组件之上构建更复杂的抽象。在合理的情况下，实现会优先考虑代码的可读性以及与标准公式一致的结构化设计，而非代码复用。

该库的初始原型曾以 jax.experimental.optix 的形式发布在 JAX 的实验性目录中。鉴于 optix 在 DeepMind 内部得到了广泛应用，并且经过多次 API 迭代后，optix 最终被移出 experimental 目录，成为独立的开源库，并更名为 optax。

Optax 的文档可以在 optax.readthedocs.io 上找到。

安装

您可以从 PyPI 安装最新发布的 Optax 版本：

pip install optax

或者，您也可以从 GitHub 安装最新的开发版本：

pip install git+https://github.com/google-deepmind/optax.git

快速入门

Optax 包含了许多流行的优化器和损失函数的实现。例如，以下代码片段使用了 optax.adam 中的 Adam 优化器以及 optax.l2_loss 中的均方误差损失。我们使用模型的参数 params 和 init 函数来初始化优化器的状态。

optimizer = optax.adam(learning_rate)
# 获取包含优化器统计信息的 `opt_state`。
params = {'w': jnp.ones((num_weights,))}
opt_state = optimizer.init(params)

为了编写更新循环，我们需要一个可以被 Jax 求导的损失函数（在此示例中使用 jax.grad）来获取梯度。

compute_loss = lambda params, x, y: optax.l2_loss(params['w'].dot(x), y)
grads = jax.grad(compute_loss)(params, xs, ys)

然后，通过 optimizer.update 将梯度转换为应应用于当前参数以获得新参数的更新值。optax.apply_updates 是一个方便的工具函数，用于执行此操作。

updates, opt_state = optimizer.update(grads, opt_state)
params = optax.apply_updates(params, updates)

您可以在 Optax 🚀 入门笔记本 中继续快速入门。

开发

我们欢迎报告问题和解决现有问题或改进功能的拉取请求。如果您有兴趣添加新功能（如新的优化器），请先打开一个问题！我们致力于使 Optax 更加灵活、通用且易于使用，以便您能够定义自己的优化器。

源代码

您可以通过以下命令查看最新源代码：

git clone https://github.com/google-deepmind/optax.git

测试

要运行测试，请执行以下脚本：

sh test.sh

文档

要构建文档，请先确保所有依赖项已安装：

pip install -e ".[docs]"

然后执行以下命令：

cd docs
make html

基准测试

一些基准测试：

• 神经网络训练算法的基准测试，Dahl G. 等，2023 年,
• 穿越拥挤的山谷——深度学习优化器的基准测试，Schmidt R. 等，2021 年.

创建您自己的基准测试：

• Benchopt：可重复、高效且协作式的优化基准测试，Moreau T. 等，2022 年。

优化器调优手册：

• 深度学习调优指南，Godbole V. 等，2023 年。

JAX 中其他与优化相关的库

• optimistix：非线性求解器，包括根查找、最小化、不动点和最小二乘法。
• matfree：无矩阵方法，可用于研究深度学习中的曲率动力学。

引用 Optax

本仓库是 DeepMind JAX 生态系统的一部分，引用 Optax 时请使用以下参考文献：

@software{deepmind2020jax,
  title = {The {D}eep{M}ind {JAX} {E}cosystem},
  author = {DeepMind and Babuschkin, Igor and Baumli, Kate and Bell, Alison and Bhupatiraju, Surya and Bruce, Jake and Buchlovsky, Peter and Budden, David and Cai, Trevor and Clark, Aidan and Danihelka, Ivo and Dedieu, Antoine and Fantacci, Claudio and Godwin, Jonathan and Jones, Chris and Hemsley, Ross and Hennigan, Tom and Hessel, Matteo and Hou, Shaobo and Kapturowski, Steven and Keck, Thomas and Kemaev, Iurii and King, Michael and Kunesch, Markus and Martens, Lena and Merzic, Hamza and Mikulik, Vladimir and Norman, Tamara and Papamakarios, George and Quan, John and Ring, Roman and Ruiz, Francisco and Sanchez, Alvaro and Sartran, Laurent and Schneider, Rosalia and Sezener, Eren and Spencer, Stephen and Srinivasan, Srivatsan and Stanojevi\'{c}, Milo\v{s} and Stokowiec, Wojciech and Wang, Luyu and Zhou, Guangyao and Viola, Fabio},
  url = {http://github.com/google-deepmind},
  year = {2020},
}

Optax 快速上手指南

Optax 是一个专为 JAX 设计的梯度处理和优化库。它提供了一系列可组合的基础构建块，帮助开发者轻松构建自定义优化器或梯度处理组件，广泛应用于深度学习研究中。

环境准备

在使用 Optax 之前，请确保您的开发环境满足以下要求：

• 操作系统：Linux、macOS 或 Windows（推荐 Linux 环境以获得最佳性能）。
• Python 版本：建议 Python 3.9 或更高版本。
• 核心依赖：
  • JAX：Optax 强依赖于 JAX 进行自动微分和数组操作。
  • NumPy / JAX NumPy。
• 硬件加速（可选但推荐）：如需使用 GPU 或 TPU 加速，请预先安装对应版本的 jaxlib（例如 pip install "jax[cuda]"）。

安装步骤

您可以通过 PyPI 安装最新稳定版，或从 GitHub 安装最新开发版。

方式一：安装稳定版（推荐）

pip install optax

国内加速提示：如果您在中国大陆地区，建议使用国内镜像源以加快下载速度：

pip install optax -i https://pypi.tuna.tsinghua.edu.cn/simple

方式二：安装开发版

如果您需要体验最新功能或参与贡献：

pip install git+https://github.com/google-deepmind/optax.git

基本使用

Optax 的使用流程通常分为三步：初始化优化器 -> 计算梯度 -> 更新参数。

以下是一个使用 Adam 优化器和均方误差损失函数的最小化示例：

import jax
import jax.numpy as jnp
import optax

# 1. 初始化优化器
# 定义学习率并创建 Adam 优化器实例
learning_rate = 0.001
optimizer = optax.adam(learning_rate)

# 模拟模型参数 (例如权重 w)
params = {'w': jnp.ones((10,))}

# 初始化优化器状态 (包含动量等统计信息)
opt_state = optimizer.init(params)

# 2. 定义损失函数并计算梯度
# 定义一个简单的损失函数：均方误差 (L2 Loss)
compute_loss = lambda params, x, y: optax.l2_loss(jnp.dot(params['w'], x), y)

# 模拟输入数据 (x) 和目标值 (y)
xs = jnp.ones((10,))
ys = jnp.array([1.0])

# 使用 jax.grad 计算梯度
grads = jax.grad(compute_loss)(params, xs, ys)

# 3. 更新参数
# 根据梯度和当前状态计算更新量
updates, opt_state = optimizer.update(grads, opt_state)

# 应用更新量到参数上，得到新参数
params = optax.apply_updates(params, updates)

print("参数更新完成，新参数:", params)

更多高级用法（如梯度裁剪、学习率调度、多优化器组合等）请参考 Optax 官方文档 或查看官方的 Getting Started Notebook。