Sonnet

文档 | 示例

Sonnet 是一个构建在 TensorFlow 2 之上的库，旨在为机器学习研究提供简单、可组合的抽象。

简介

Sonnet 由 DeepMind 的研究人员设计和开发。它可以用于构建各种用途的神经网络（无监督/有监督学习、强化学习等）。我们发现它对我们的团队非常有效，你也许也会喜欢！

具体来说，Sonnet 提供了一个简单而强大的编程模型，核心概念是 snt.Module。模块可以包含参数、其他模块以及对用户输入应用某些函数的方法。Sonnet 自带许多预定义模块（如 snt.Linear、snt.Conv2D、snt.BatchNorm）和一些预定义的模块网络（如 snt.nets.MLP），同时也鼓励用户自定义模块。

与许多框架不同，Sonnet 对 如何 使用模块持开放态度。模块被设计成自包含且彼此完全解耦。Sonnet 不自带训练框架，而是鼓励用户自行构建或采用他人开发的框架。

Sonnet 还致力于简洁易懂，我们的代码（希望）清晰且专注。当我们选择默认值时（例如初始参数值的默认值），我们会说明原因。

快速入门

示例

尝试 Sonnet 的最简单方式是使用 Google Colab，它提供免费的 Python 笔记本，并配备 GPU 或 TPU。

• 使用 MLP 预测 MNIST
• 在 MNIST 上训练一个小 GAN
• 使用 snt.distribute 进行分布式训练

安装

要开始使用，请安装 TensorFlow 2.0 和 Sonnet 2：

$ pip install tensorflow tensorflow-probability
$ pip install dm-sonnet

你可以运行以下代码来验证是否安装正确：

import tensorflow as tf
import sonnet as snt

print("TensorFlow 版本 {}".format(tf.__version__))
print("Sonnet 版本 {}".format(snt.__version__))

使用现有模块

Sonnet 自带许多内置模块，你可以直接使用。例如，要定义一个 MLP，我们可以使用 snt.Sequential 模块来按顺序调用多个模块，将前一个模块的输出作为下一个模块的输入。我们可以使用 snt.Linear 和 tf.nn.relu 来实际定义计算过程：

mlp = snt.Sequential([
    snt.Linear(1024),
    tf.nn.relu,
    snt.Linear(10),
])

要使用这个模块，我们需要“调用”它。Sequential 模块（以及大多数模块）都定义了 __call__ 方法，这意味着你可以直接通过名称调用它们：

logits = mlp(tf.random.normal([batch_size, input_size]))

通常，我们还会获取模块的所有参数。Sonnet 中的大多数模块会在第一次被调用并传入输入时创建其参数（因为参数的形状通常是输入的函数）。Sonnet 模块提供了两个属性来访问这些参数。

variables 属性会返回给定模块引用的所有 tf.Variable：

all_variables = mlp.variables

需要注意的是，tf.Variable 并不只用于模型的参数。例如，它们还用于存储 snt.BatchNorm 中指标的状态。大多数情况下，用户会提取模块的变量以传递给优化器进行更新。然而，不可训练的变量通常不应包含在该列表中，因为它们是通过另一种机制更新的。TensorFlow 内置了标记变量为“可训练”（模型参数）或“不可训练”（其他变量）的机制。Sonnet 提供了一种从模块中收集所有可训练变量的方法，这通常是你要传递给优化器的内容：

model_parameters = mlp.trainable_variables

构建自己的模块

Sonnet 强烈鼓励用户继承 snt.Module 来定义自己的模块。让我们从创建一个简单的线性层 MyLinear 开始：

class MyLinear(snt.Module):

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

  @snt.once
  def _initialize(self, x):
    initial_w = tf.random.normal([x.shape[1], self.output_size])
    self.w = tf.Variable(initial_w, name="w")
    self.b = tf.Variable(tf.zeros([self.output_size]), name="b")

  def __call__(self, x):
    self._initialize(x)
    return tf.matmul(x, self.w) + self.b

使用这个模块非常简单：

mod = MyLinear(32)
mod(tf.ones([batch_size, input_size]))

通过继承 snt.Module，你可以免费获得许多便利的功能。例如，默认实现的 __repr__ 方法，会显示构造函数的参数（非常有助于调试和内省）：

>>> print(repr(mod))
MyLinear(output_size=10)

你还可以使用 variables 和 trainable_variables 属性：

>>> mod.variables
(<tf.Variable 'my_linear/b:0' shape=(10,) ...>,
 <tf.Variable 'my_linear/w:0' shape=(1, 10) ...>)

你可能会注意到上述变量名前都有 my_linear 前缀。这是因为当方法被调用时，Sonnet 模块会进入模块的命名空间。通过进入模块的命名空间，我们可以为 TensorBoard 等工具提供更有用的图结构（例如，所有发生在 my_linear 内的操作都会被归类到名为 my_linear 的组中）。

此外，你的模块现在也支持 TensorFlow 的检查点保存和模型保存功能，这些高级特性将在后面介绍。

序列化

Sonnet 支持多种序列化格式。我们支持的最简单格式是 Python 的 pickle，所有内置模块都经过测试，确保可以在同一 Python 进程中通过 pickle 保存和加载。不过，我们通常不建议使用 pickle，因为它在 TensorFlow 的许多部分中支持不佳，而且在实践中往往比较脆弱。

TensorFlow 检查点

参考： https://www.tensorflow.org/alpha/guide/checkpoints

TensorFlow 检查点可用于在训练过程中定期保存参数值。这在程序崩溃或中断时保存训练进度非常有用。Sonnet 与 TensorFlow 检查点能够很好地协同工作：

checkpoint_root = "/tmp/checkpoints"
checkpoint_name = "example"
save_prefix = os.path.join(checkpoint_root, checkpoint_name)

my_module = create_my_sonnet_module()  # 可以是任何继承自 snt.Module 的对象。

# `Checkpoint` 对象管理与其构造函数中传入的对象相关联的 TensorFlow 状态的检查点保存。请注意，Checkpoint 支持在创建时恢复，这意味着 `my_module` 的变量**无需**在从检查点恢复之前被创建（它们的值将在被创建时被恢复）。
checkpoint = tf.train.Checkpoint(module=my_module)

# 大多数训练脚本都会希望在存在检查点的情况下从中恢复。例如，如果你中断了训练（比如为了将 GPU 用于其他任务，或者在云环境中实例被抢占），就需要进行恢复。
latest = tf.train.latest_checkpoint(checkpoint_root)
if latest is not None:
  checkpoint.restore(latest)

for step_num in range(num_steps):
  train(my_module)

  # 在训练过程中，我们会偶尔保存权重的值。需要注意的是，这是一个阻塞调用，可能会比较慢（通常是在机器上速度最慢的存储设备上写入数据）。如果你的环境更加可靠，可能可以适当减少保存频率。
  if step_num > 0 and step_num % 1000 == 0:
    checkpoint.save(save_prefix)

# 务必保存最终的模型参数！！
checkpoint.save(save_prefix)

TensorFlow Saved Model

参考: https://www.tensorflow.org/alpha/guide/saved_model

TensorFlow 的 Saved Model 可以用来保存一份与 Python 源代码解耦的网络副本。这是通过保存描述计算过程的 TensorFlow 图以及包含权重值的检查点来实现的。

要创建一个 Saved Model，首先需要创建一个你想要保存的 snt.Module：

my_module = snt.nets.MLP([1024, 1024, 10])
my_module(tf.ones([1, input_size]))

接下来，我们需要创建另一个模块，用于描述我们想要导出的模型的具体部分。我们建议这样做（而不是直接修改原始模型），以便对实际导出的内容有更精细的控制。这通常很重要，可以避免生成非常大的 Saved Model，并且只共享你希望公开的部分（例如，你只想共享 GAN 的生成器，而保持判别器私密）。

@tf.function(input_signature=[tf.TensorSpec([None, input_size])])
def inference(x):
  return my_module(x)

to_save = snt.Module()
to_save.inference = inference
to_save.all_variables = list(my_module.variables)
tf.saved_model.save(to_save, "/tmp/example_saved_model")

现在，我们在 /tmp/example_saved_model 文件夹中有了一个 Saved Model：

$ ls -lh /tmp/example_saved_model
total 24K
drwxrwsr-t 2 tomhennigan 154432098 4.0K Apr 28 00:14 assets
-rw-rw-r-- 1 tomhennigan 154432098  14K Apr 28 00:15 saved_model.pb
drwxrwsr-t 2 tomhennigan 154432098 4.0K Apr 28 00:15 variables

加载这个模型非常简单，甚至可以在没有构建该 Saved Model 的任何 Python 代码的另一台机器上完成：

loaded = tf.saved_model.load("/tmp/example_saved_model")

# 使用 inference 方法。请注意，这并不会运行 `to_save` 中的 Python 代码，而是使用 Saved Model 中包含的 TensorFlow 图。
loaded.inference(tf.ones([1, input_size]))

# 可以使用 all_variables 属性来获取已恢复的变量。
assert len(loaded.all_variables) > 0

需要注意的是，加载后的对象并不是 Sonnet 模块，而是一个容器对象，它包含了我们在前面步骤中添加的特定方法（如 inference）和属性（如 all_variables）。

分布式训练

示例: https://github.com/deepmind/sonnet/blob/v2/examples/distributed_cifar10.ipynb

Sonnet 支持使用 自定义的 TensorFlow 分布式策略 进行分布式训练。

Sonnet 与使用 tf.keras 进行分布式训练的一个关键区别在于，Sonnet 模块和优化器在分布式策略下运行时的行为不会发生变化（例如，我们不会对梯度进行平均或同步批归一化统计信息）。我们认为用户应该完全掌控这些训练细节，而不应将其硬编码到库中。相应的权衡是，你需要在自己的训练脚本中实现这些功能（通常只需两行代码，在应用优化器之前对梯度进行全归约操作），或者替换为明确支持分布式训练的模块（例如 snt.distribute.CrossReplicaBatchNorm）。

我们的 分布式 Cifar-10 示例 详细介绍了如何使用 Sonnet 进行多 GPU 训练。

Sonnet 快速上手指南

Sonnet 是由 DeepMind 基于 TensorFlow 2 构建的神经网络库，旨在为机器学习研究提供简单、可组合的抽象模块（snt.Module）。它不强制绑定训练框架，赋予开发者最大的灵活性。

环境准备

• 系统要求：支持 Linux、macOS 或 Windows。
• 前置依赖：
  • Python 3.7+
  • TensorFlow 2.x
  • TensorFlow Probability (可选，部分功能需要)
• 硬件建议：推荐使用配备 GPU 或 TPU 的环境以获得最佳性能（如使用 Google Colab 或本地 CUDA 环境）。

国内加速提示：由于官方源位于海外，建议使用国内镜像源安装依赖以提升速度。

• 清华源：https://pypi.tuna.tsinghua.edu.cn/simple
• 阿里源：https://mirrors.aliyun.com/pypi/simple/

安装步骤

使用 pip 安装 TensorFlow 和 Sonnet。推荐在虚拟环境中操作。

1. 安装依赖

# 使用国内镜像源加速安装
pip install -i https://pypi.tuna.tsinghua.edu.cn/simple tensorflow tensorflow-probability
pip install -i https://pypi.tuna.tsinghua.edu.cn/simple dm-sonnet

2. 验证安装

运行以下 Python 代码确认安装成功：

import tensorflow as tf
import sonnet as snt

print("TensorFlow version {}".format(tf.__version__))
print("Sonnet version {}".format(snt.__version__))

基本使用

Sonnet 的核心概念是 snt.Module。你可以直接使用内置模块，也可以继承 snt.Module 创建自定义模块。

1. 使用内置模块构建 MLP

Sonnet 提供了丰富的预定义模块（如 snt.Linear, snt.Conv2D）。利用 snt.Sequential 可以轻松堆叠网络层。

import tensorflow as tf
import sonnet as snt

# 定义一个多层感知机 (MLP)
mlp = snt.Sequential([
    snt.Linear(1024),
    tf.nn.relu,
    snt.Linear(10),
])

# 调用模块进行前向传播
# 参数会在第一次调用时根据输入形状自动初始化
batch_size = 32
input_size = 784
logits = mlp(tf.random.normal([batch_size, input_size]))

# 获取可训练参数（用于优化器）
model_parameters = mlp.trainable_variables

2. 构建自定义模块

通过继承 snt.Module 类，你可以定义自己的网络层。Sonnet 会自动处理变量作用域、序列化及调试信息展示。

class MyLinear(snt.Module):
    def __init__(self, output_size, name=None):
        super(MyLinear, self).__init__(name=name)
        self.output_size = output_size

    @snt.once
    def _initialize(self, x):
        # 延迟初始化：根据输入形状创建变量
        initial_w = tf.random.normal([x.shape[1], self.output_size])
        self.w = tf.Variable(initial_w, name="w")
        self.b = tf.Variable(tf.zeros([self.output_size]), name="b")

    def __call__(self, x):
        self._initialize(x)
        return tf.matmul(x, self.w) + self.b

# 使用自定义模块
mod = MyLinear(32)
output = mod(tf.ones([batch_size, input_size]))

# 查看模块表示和变量
print(repr(mod)) 
print(mod.variables)

下一步

• 示例代码：访问 GitHub Examples 查看 MNIST 分类、GAN 训练及分布式训练的完整 Notebook。
• 文档：详细 API 参考请参阅 官方文档。