Optax

CI状态 pypi

简介

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

Optax的设计旨在通过提供可以轻松以自定义方式重新组合的构建模块来促进研究。

我们的目标是：

提供核心组件的简单、经过充分测试和高效的实现。
通过使研究人员能够轻松地将低级组件组合成自定义优化器（或其他梯度处理组件），提高研究生产力。
通过让任何人都能轻松贡献，加速新想法的采用。

我们倾向于专注于可以有效组合成自定义解决方案的小型可组合构建模块。其他人可以在这些基本组件的基础上构建更复杂的抽象。在合理的情况下，实现优先考虑可读性和将代码结构与标准方程匹配，而不是代码重用。

该库的初始原型在JAX的实验文件夹中作为jax.experimental.optix提供。鉴于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的均方误差。我们使用init函数和模型的params初始化优化器状态。

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 🚀 入门笔记本中继续快速入门。

开发

我们欢迎新的贡献者。

源代码

您可以使用以下命令检查最新的源代码。

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

测试

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

sh ./test.sh

文档

要构建文档，首先确保安装了所有依赖项。

pip install -e ".[docs]"

然后，执行以下操作。

cd docs/
make html

基准测试

如果您在众多可用的深度学习优化器中感到迷失，有一些广泛的基准测试：

神经网络训练算法基准测试，Dahl G. et al, 2023,

通过拥挤的谷底下降 — 深度学习优化器基准测试，Schmidt R. et al, 2021。

如果您有兴趣为某些任务开发自己的基准测试，请考虑以下框架

Benchopt：可重现、高效和协作的优化基准测试，Moreau T. et al, 2022。

最后，如果您正在寻找一些关于调优优化器的建议，请考虑查看

深度学习调优手册，Godbole V. et al, 2023。

引用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},
}