Jax 生态再添新库：DeepMind 开源 Haiku、RLax

机器之心报道

Jax 是一个优秀的代码库，在进行科学计算的同时能够自动微分，还有 GPU、TPU 的性能加速加持。但是 Jax 的生态还不够完善，使用者相比 TF、PyTorch 少得多。近日，DeepMind 开源了两个基于 Jax 的新库，给这个生态注入了新的活力。

Jax 是谷歌开源的一个科学计算库，能对 Python 程序与 NumPy 运算执行自动微分，而且能够在 GPU 和 TPU 上运行，具有很高的性能。基于 Jax 已有很多优秀的开源项目，如 Trax 等。近日，DeepMind 开源了两个基于 Jax 的新机器学习库，分别是 Haiku 和 RLax，它们都有着各自的特色，对于丰富深度学习社区框架、提升研究者和开发者的使用体验有着不小的意义。

Haiku：https://github.com/deepmind/haiku

RLax：https://github.com/deepmind/rlax

首先值得注意的是 Haiku，这是一个面向 Jax 的深度学习代码库，它是由 Sonnet 作者——一个谷歌的神经网络库团队开发的。

为什么要使用 Haiku？这是因为其支持的是 Jax，Jax 在灵活性和性能上具有相当的优势。但是另一方面，Jax 本身是函数式的，和面向对象的用户习惯有差别。因此，通过 Haiku，用户可以在 Jax 上进行面向对象开发了。

此外，Haiku 的 API 和编程模型都是基于 Sonnet，因此使用过 Sonnet 的用户可以快速上手。项目作者也表示，Sonnet 之于 TensorFlow 的提升就如同 Haiku 之于 Jax。

目前，Haiku 已公开了 Alpha 版本，已完全开源。项目作者欢迎使用者提出建议。

Haiku 主要分为两个模块：hk.Modules和hk.transform。下文将会分别介绍。

hk.Modules是 Python 对象，保存着到参数、其他模块和方法的参照（references）。

hk.transform 则负责将面向对象的模块转换为纯粹的函数式代码，然后让 jax 中的 jax.jit, jax.grad, jax.pmap 等进行处理，从而实现和 Jax 组件的兼容。

Haiku 能够做到很多机器学习需要完成的任务，相关功能和代码如下：

在 Haiku 中，类似于 TF2.0 和 PyTorch，你可以自定义模块，作为 hk.Module 的子类。例如，自定义一个线性层：

可以看出，Haiku 的代码和 TensorFlow 等非常相似，但是你可以看到包括 numpy 等的方法还可以定义在模块中。Haiku 的优势就在于，它不是一个封闭的框架，而是代码库，因此可以在定义模块的过程中调用其他的库和方法。

当定义好线性层后，我们想要试试自动微分的方法了：

这里可以看到，定义好模块和前向传播的函数后，使用 hk.transform(forward_fn) 可以将这种面向对象的方法转换成 Jax 底层的函数式代码进行处理，因此你不需要担心底层的计算问题。另外，这里的代码相比 TensorFlow 还要简洁。

有时候，我们想要在训练的过程中保持某些内部参数的状态，在 Haiku 上这也是非常容易实现的。

如上所示，只需要两行代码进行设置。

由于所有的代码都会被转换成 Jax 的函数，因此它们和 jax.pmap. 是完全兼容的。这说明，我们可以利用 jax.pmap 来进行分布式计算。

如下为进行数据分割的分布式加速代码，首先，我们先定义模型和训练步骤：

然后设定将参数拷贝到所有的设备上：

定义数据分批的方法，以及参数更新的方法：

最后开始分布式计算即可：

除了令人印象深刻的 Haiku 外，DeepMind 还开源了 RLax——这是一个基于 Jax 的强化学习库。

相比 Haiku，RLax 专门针对强化学习。项目作者认为，尽管强化学习中的算子和函数并不是完全的算法，但是，如果需要构建完全基于函数式的智能体，就需要特定的数学算子。

因此，函数式的 Jax 就成为了一个不错的选择。在 Jax 上进行一定的开发后，就可以有专用的强化学习库了。RLax 目前的资料还较少，但项目已提供了一个示例代码：使用 RLax 进行 Q-learning 模型的搭建和训练。

代码如下，首先，使用 Haiku 构建基本的强化学习模型：

设定训练的方法：

以下和 Jax 结合，定义策略、奖励等：

可以看到，RLax 基于 jax.jit 的方法，在性能方面有不错的提升。更有趣的是，构建模型的过程中使用了前文提到的 Haiku，可见基于 Jax 生态的代码库之间都是可以兼容的。

从 DeepMind 近日开源的两个代码库可以看到，虽然现在深度学习框架依然在稳步发展，但是针对高性能的科学计算也渐渐变得更为重要了。而 Jax 这样的优秀开源项目，无疑也需要更多的生态支持。这次开源的 Haiku 和 RLax，无疑能够巩固 Jax 的地位，使其优秀的特性进一步得到发挥。

------------------------------------------------