俳句：{JAX}的{十四行诗}

概述 | 为什么选择俳句? | 快速开始 | 安装 | 示例 | 用户手册 | 文档 | 引用俳句

pytest docs pypi

[!重要] 📣 截至2023年7月，Google DeepMind建议新项目采用Flax而不是俳句。Flax是由Google Brain最初开发，现在由Google DeepMind开发的神经网络库。 📣

在撰写本文时，Flax拥有比俳句更多的特性，较大的开发团队和更活跃的社区。Flax在用户群体中有更高的接受度，具有更广泛的文档，示例和一个活跃的社区来创建从头到尾的示例。

俳句将保持尽力支持，但项目将进入维护模式，这意味着开发工作的重点将放在错误修复和与JAX新版本的兼容性上。

为了保持俳句与Python和JAX的新版本的兼容，我们将继续发布新的版本，但不会添加（或接受PR）新功能。

我们在Google DeepMind内部广泛使用俳句，并计划无限期地以这种模式支持俳句。

什么是俳句？

俳句是一种工具<br> 用于构建神经网络<br> 想象一下：“{JAX}的{十四行诗}”

俳句是由Sonnet的一些作者为JAX开发的简单神经网络库，Sonnet是一个为TensorFlow开发的神经网络库。

俳句的文档可以在https://dm-haiku.readthedocs.io/找到。

**说明：**如果您在寻找操作系统俳句，请访问https://haiku-os.org/。

概述<a id="overview"></a>

JAX是一个结合NumPy、自动微分和一流GPU/TPU支持的数值计算库。

俳句是一个简单的JAX神经网络库，允许用户使用熟悉的面向对象编程模型，同时允许完全访问JAX的纯函数转换。

俳句提供了两个核心工具：模块抽象hk.Module和一个简单的函数转换hk.transform。

hk.Module是持有自身参数、其他模块和应用用户输入函数的方法的Python对象。

hk.transform将使用这些面向对象、功能上“纯净”模块的函数转换成可以与jax.jit、jax.grad、jax.pmap等一起使用的纯函数。

为什么选择俳句？<a id="why-haiku"></a>

有很多为JAX开发的神经网络库。为什么要选择俳句？

俳句已经通过DeepMind的研究人员在大规模测试。

DeepMind已经用俳句和JAX重新实现了许多实验的结果，包括图像和语言处理、大规模的生成模型和强化学习。

俳句是一个库，而不是框架。

俳句旨在简化特定任务：管理模型参数和其他模型状态。
你可以期望俳句与其他库组合，并且与JAX的其余部分很好地协作。
除此之外，俳句设计得很少干涉用户的工作——它不定义自定义优化器、检查点格式或复制API。

俳句没有重复发明轮子。

俳句基于DeepMind几乎普遍采用的神经网络库Sonnet的编程模型和API。它保留了Sonnet的用于状态管理的Module(模块)编程模型，同时保留了对JAX的函数转换的访问。
俳句的API和抽象与Sonnet尽可能接近。许多用户发现Sonnet是TensorFlow中一个富有成效的编程模型；俳句在JAX中实现了相同的体验。

迁移到俳句很容易。

通过设计，从TensorFlow和Sonnet迁移到JAX和俳句是很容易的。
除了新功能（如hk.transform）外，俳句旨在匹配Sonnet 2的API。模块、方法、参数名称、默认值和初始化方案应该匹配。

俳句简化了JAX的其他方面。

俳句提供了一个简单的模型来处理随机数。在转换后的函数中，hk.next_rng_key()返回一个唯一的随机数生成器键。
这些唯一键是从传递给顶级转换函数的初始随机键派生的，因此可以安全地与JAX程序转换一起使用。

快速开始<a id="quickstart"></a>

让我们来看一个示例神经网络、损失函数和训练循环。(欲了解更多示例, 请参见我们的示例目录。 MNIST 示例是一个很好的起点。）

import haiku as hk
import jax.numpy as jnp

def softmax_cross_entropy(logits, labels):
  one_hot = jax.nn.one_hot(labels, logits.shape[-1])
  return -jnp.sum(jax.nn.log_softmax(logits) * one_hot, axis=-1)

def loss_fn(images, labels):
  mlp = hk.Sequential([
      hk.Linear(300), jax.nn.relu,
      hk.Linear(100), jax.nn.relu,
      hk.Linear(10),
  ])
  logits = mlp(images)
  return jnp.mean(softmax_cross_entropy(logits, labels))

loss_fn_t = hk.transform(loss_fn)
loss_fn_t = hk.without_apply_rng(loss_fn_t)

rng = jax.random.PRNGKey(42)
dummy_images, dummy_labels = next(input_dataset)
params = loss_fn_t.init(rng, dummy_images, dummy_labels)

def update_rule(param, update):
  return param - 0.01 * update

for images, labels in input_dataset:
  grads = jax.grad(loss_fn_t.apply)(params, images, labels)
  params = jax.tree_util.tree_map(update_rule, params, grads)

俳句的核心是hk.transform。transform函数允许您编写依赖于参数的神经网络函数（此处是Linear层的权重）而不需要明确编写初始化那些参数的样板代码。transform通过将函数转换成纯函数对init和apply实现的形式来做到这点。

`init`

init函数，签名为params = init(rng, ...)（其中...是未转换函数的参数），允许您收集网络中任何参数的初始值。俳句通过运行您的函数，跟踪任何通过hk.get_parameter（由如hk.Linear调用）请求的参数并返回给您。

params对象是您的网络中所有参数的嵌套数据结构，设计供您检查和操作。具体地，它是模块名称到模块参数的映射，其中模块参数是参数名称到参数值的映射。例如：

{'linear': {'b': ndarray(..., shape=(300,), dtype=float32),
            'w': ndarray(..., shape=(28, 300), dtype=float32)},
 'linear_1': {'b': ndarray(..., shape=(100,), dtype=float32),
              'w': ndarray(..., shape=(1000, 100), dtype=float32)},
 'linear_2': {'b': ndarray(..., shape=(10,), dtype=float32),
              'w': ndarray(..., shape=(100, 10), dtype=float32)}}

`apply`

apply函数，签名为result = apply(params, rng, ...)，允许您向函数注入参数值。当调用hk.get_parameter时，返回的值将来自您提供给apply作为输入的params：

loss = loss_fn_t.apply(params, rng, images, labels)

注意，由于我们的损失函数执行的实际计算不依赖于随机数，因此传入一个随机数生成器是没有必要的，因此我们也可以传入None给rng参数。（请注意，如果您的计算确实使用了随机数，传入None给rng将引发错误）。在上面的示例中，我们让俳句自动对我们执行此操作：

loss_fn_t = hk.without_apply_rng(loss_fn_t)

既然apply是一个纯函数，我们可以将其传递给jax.grad（或其他JAX变换）：

grads = jax.grad(loss_fn_t.apply)(params, images, labels)

训练

此示例中的训练循环非常简单。一个需要注意的细节是使用jax.tree_util.tree_map来将sgd函数应用于params和grads中的所有匹配条目。结果具有与原始params相同的结构，并可以再次用于apply。

安装<a id="installation"></a>

俳句是用纯Python编写的，但依赖于通过JAX的C++代码。

由于JAX的安装因CUDA版本而异，俳句没有在requirements.txt中列出JAX作为依赖项。

首先，按照这些说明来安装带有相关加速器支持的JAX。

然后，使用pip安装俳句：

$ pip install git+https://github.com/deepmind/dm-haiku

或者，您可以通过PyPI安装：

$ pip install -U dm-haiku

我们的示例依赖于额外的库（例如bsuite）。您可以使用pip安装所有额外的依赖：

$ pip install -r examples/requirements.txt

用户手册<a id="user-manual"></a>

编写您自己的模块

在俳句中，所有模块都是hk.Module的子类。您可以实现任何方法（没有被特殊处理），但通常模块实现__init__和__call__。

让我们一起实现一个线性层：

class MyLinear(hk.Module):

  def __init__(self, output_size, name=None):
    super().__init__(name=name)
    self.output_size = output_size

  def __call__(self, x):
    j, k = x.shape[-1], self.output_size
    w_init = hk.initializers.TruncatedNormal(1. / np.sqrt(j))
    w = hk.get_parameter("w", shape=[j, k], dtype=x.dtype, init=w_init)
    b = hk.get_parameter("b", shape=[k], dtype=x.dtype, init=jnp.zeros)
    return jnp.dot(x, w) + b

所有模块都有一个名称。如果没有传递name参数，模块名将从Python类名中推断（例如MyLinear变为my_linear）。模块可以有通过hk.get_parameter(param_name, ...)访问的命名参数。我们使用这个API（而不是直接使用对象属性），这样我们可以使用hk.transform将您的代码转换为纯函数。

使用模块时，您需要定义函数并使用hk.transform将其转换为一对纯函数。有关转换后函数返回的函数的详细信息，请参见我们的快速开始：

def forward_fn(x):
  model = MyLinear(10)
  return model(x)

# 将`forward_fn`转换为具有`init`和`apply`方法的对象。 默认情况下，
# `apply`将需要一个rng（可以是None），用于
# `hk.next_rng_key`。
forward = hk.transform(forward_fn)

x = jnp.ones([1, 1])

# 当我们运行`forward.init`时，俳句将运行`forward_fn(x)`并收集初始参数值。由于参数
# 通常是随机初始化的，因此俳句要求您传递一个随机数生成器键给`init`：
key = hk.PRNGSequence(42)
params = forward.init(next(key), x)

# 当我们运行`forward.apply`时，俳句将运行`forward_fn(x)`并从作为第一个参数传递的`params`中注入参数
# 值。 请注意，通过`hk.transform(f)`转换的模型必须使用额外的
# `rng`参数调用：`forward.apply(params, rng, x)`。 使用
# `hk.without_apply_rng(hk.transform(f))`如果不需要这样做。
y = forward.apply(params, None, x)

处理随机模型

一些模型可能需要随机采样作为计算的一部分。例如，在使用重参数化技巧的变分自编码器中，需要从标准正态分布中进行随机采样。对于dropout，我们需要一个随机掩码来从输入中丢弃单元。使其与JAX一起工作的主要障碍在于PRNG键的管理。

在Haiku中，我们提供了一个简单的API来维护与模块关联的PRNG键序列：hk.next_rng_key()（或对于多个键使用next_rng_keys()）：

class MyDropout(hk.Module):

  def __init__(self, rate=0.5, name=None):
    super().__init__(name=name)
    self.rate = rate

  def __call__(self, x):
    key = hk.next_rng_key()
    p = jax.random.bernoulli(key, 1.0 - self.rate, shape=x.shape)
    return x * p / (1.0 - self.rate)

forward = hk.transform(lambda x: MyDropout()(x))

key1, key2 = jax.random.split(jax.random.PRNGKey(42), 2)
params = forward.init(key1, x)
prediction = forward.apply(params, key2, x)

要更全面地了解与随机模型一起工作的情况，请参阅我们的VAE示例。

注意： hk.next_rng_key()不是功能纯的，这意味着你应该避免在hk.transform内使用它与JAX变换一起使用。欲了解更多信息和可能的解决方法，请查阅Haiku变换的文档和可用的Haiku网络中的JAX变换包装器。

使用不可训练状态

一些模型可能希望维护一些内部的、可变的状态。例如，在批量归一化中，训练过程中遇到的值的移动平均值是维护的。

在Haiku中，我们提供了一个简单的API来维护与模块关联的可变状态：hk.set_state和hk.get_state。使用这些函数时，您需要使用hk.transform_with_state转换您的函数，因为返回的函数对的签名是不同的：

def forward(x, is_training):
  net = hk.nets.ResNet50(1000)
  return net(x, is_training)

forward = hk.transform_with_state(forward)

# `init`函数现在返回参数和状态。状态包含使用`hk.set_state`创建的任何内容。结构与参数相同（例如，这是一个按模块命名值的映射）。
params, state = forward.init(rng, x, is_training=True)

# `apply`函数现在接受参数和状态。此外，它将返回更新的状态值。在resnet示例中，这将是用于批量规范化层中的移动平均值的更新值。
logits, state = forward.apply(params, state, rng, x, is_training=True)

如果你忘记使用hk.transform_with_state，不要担心，我们会打印一个明确的错误，指向你hk.transform_with_state，而不是默默地丢弃你的状态。

使用`jax.pmap`进行分布式训练

从hk.transform（或hk.transform_with_state）返回的纯函数完全兼容jax.pmap。有关使用jax.pmap进行SPMD编程的更多详细信息，请查看此处。

在Haiku中使用jax.pmap的一个常见用途是对许多加速器进行数据并行训练，可能跨多个主机。在Haiku中，这可能如下所示：

def loss_fn(inputs, labels):
  logits = hk.nets.MLP([8, 4, 2])(x)
  return jnp.mean(softmax_cross_entropy(logits, labels))

loss_fn_t = hk.transform(loss_fn)
loss_fn_t = hk.without_apply_rng(loss_fn_t)

# 在单个设备上初始化模型。
rng = jax.random.PRNGKey(428)
sample_image, sample_label = next(input_dataset)
params = loss_fn_t.init(rng, sample_image, sample_label)

# 将参数复制到所有设备上。
num_devices = jax.local_device_count()
params = jax.tree_util.tree_map(lambda x: np.stack([x] * num_devices), params)

def make_superbatch():
  """构造一个超级批次，即每个设备一个数据批次"""
  # 获取N个批次，然后拆分成图像列表和标签列表。
  superbatch = [next(input_dataset) for _ in range(num_devices)]
  superbatch_images, superbatch_labels = zip(*superbatch)
  # 将超级批次堆叠为一个具有前导维度的数组，而不是一个Python列表。这是`jax.pmap`期望的输入。
  superbatch_images = np.stack(superbatch_images)
  superbatch_labels = np.stack(superbatch_labels)
  return superbatch_images, superbatch_labels

def update(params, inputs, labels, axis_name='i'):
  """基于输入和标签的表现更新参数。"""
  grads = jax.grad(loss_fn_t.apply)(params, inputs, labels)
  # 在所有数据并行副本之间取梯度的平均值。
  grads = jax.lax.pmean(grads, axis_name)
  # 使用SGD或Adam或...更新参数
  new_params = my_update_rule(params, grads)
  return new_params

# 进行几次训练更新。
for _ in range(10):
  superbatch_images, superbatch_labels = make_superbatch()
  params = jax.pmap(update, axis_name='i')(params, superbatch_images,
                                           superbatch_labels)

要更全面地了解分布式Haiku训练，请查看我们的ImageNet上的ResNet-50示例。

引用Haiku<a id="citing-haiku"></a>

要引用此存储库：

@software{haiku2020github,
  author = {Tom Hennigan and Trevor Cai and Tamara Norman and Lena Martens and Igor Babuschkin},
  title = {{H}aiku: {S}onnet for {JAX}},
  url = {http://github.com/deepmind/dm-haiku},
  version = {0.0.10},
  year = {2020},
}

在此bibtex条目中，版本号应取自haiku/__init__.py，年份对应于项目的开源发布年份。