训练检查点

在 TensorFlow.org 上查看 在 Google Colab 中运行 在 GitHub 上查看源码 下载笔记本

短语“保存 TensorFlow 模型”通常有两种含义：

  1. 检查点 (Checkpoints)，或
  2. SavedModel。

检查点捕获模型所使用的所有参数（tf.Variable 对象）的精确值。检查点不包含模型所定义计算的任何描述，因此通常仅在可以使用已保存参数值的源代码可用时才有用。

另一方面，SavedModel 格式除了包含参数值（检查点）外，还包括模型所定义计算的序列化描述。此格式的模型独立于创建模型的源代码。因此，它们适用于通过 TensorFlow Serving、TensorFlow Lite、TensorFlow.js 或其他编程语言（C、C++、Java、Go、Rust、C# 等 TensorFlow API）的程序进行部署。

本指南涵盖用于写入和读取检查点的 API。

设置

import tensorflow as tf

class Net(tf.keras.Model):
  """A simple linear model."""

  def __init__(self):
    super(Net, self).__init__()
    self.l1 = tf.keras.layers.Dense(5)

  def call(self, x):
    return self.l1(x)

net = Net()

从 tf.keras 训练 API 保存

请参阅关于保存和恢复的 tf.keras 指南。

tf.keras.Model.save_weights 会保存一个 TensorFlow 检查点。

net.save_weights('easy_checkpoint')

写入检查点

TensorFlow 模型的持久状态存储在 tf.Variable 对象中。这些对象可以直接构建，但通常通过高级 API（如 tf.keras.layers 或 tf.keras.Model）创建。

管理变量的最简单方法是将它们附加到 Python 对象，然后引用这些对象。

tf.train.Checkpoint、tf.keras.layers.Layer 和 tf.keras.Model 的子类会自动跟踪分配给其属性的变量。以下示例构建了一个简单的线性模型，然后写入包含模型所有变量值的检查点。

您可以使用 Model.save_weights 轻松保存模型检查点。

手动设置检查点

设置

为了演示 tf.train.Checkpoint 的所有功能，请定义一个玩具数据集和优化步骤。 

def toy_dataset():
  inputs = tf.range(10.)[:, None]
  labels = inputs * 5. + tf.range(5.)[None, :]
  return tf.data.Dataset.from_tensor_slices(
    dict(x=inputs, y=labels)).repeat().batch(2)

def train_step(net, example, optimizer):
  """Trains `net` on `example` using `optimizer`."""
  with tf.GradientTape() as tape:
    output = net(example['x'])
    loss = tf.reduce_mean(tf.abs(output - example['y']))
  variables = net.trainable_variables
  gradients = tape.gradient(loss, variables)
  optimizer.apply_gradients(zip(gradients, variables))
  return loss

创建检查点对象

使用 tf.train.Checkpoint 对象手动创建检查点，将您想要保存检查点的对象设置为该对象的属性。

tf.train.CheckpointManager 也可以用于管理多个检查点。 

opt = tf.keras.optimizers.Adam(0.1)
dataset = toy_dataset()
iterator = iter(dataset)
ckpt = tf.train.Checkpoint(step=tf.Variable(1), optimizer=opt, net=net, iterator=iterator)
manager = tf.train.CheckpointManager(ckpt, './tf_ckpts', max_to_keep=3)

训练并保存模型检查点

下面的训练循环创建了一个模型实例和一个优化器实例，然后将它们收集到一个 tf.train.Checkpoint 对象中。它在循环中对每个数据批次调用训练步骤，并定期将检查点写入磁盘。 

def train_and_checkpoint(net, manager):
  ckpt.restore(manager.latest_checkpoint)
  if manager.latest_checkpoint:
    print("Restored from {}".format(manager.latest_checkpoint))
  else:
    print("Initializing from scratch.")

  for _ in range(50):
    example = next(iterator)
    loss = train_step(net, example, opt)
    ckpt.step.assign_add(1)
    if int(ckpt.step) % 10 == 0:
      save_path = manager.save()
      print("Saved checkpoint for step {}: {}".format(int(ckpt.step), save_path))
      print("loss {:1.2f}".format(loss.numpy()))

train_and_checkpoint(net, manager)

恢复并继续训练

在第一个训练周期之后，您可以传递一个新的模型和管理器，但可以从中断的地方准确地继续训练。 

opt = tf.keras.optimizers.Adam(0.1)
net = Net()
dataset = toy_dataset()
iterator = iter(dataset)
ckpt = tf.train.Checkpoint(step=tf.Variable(1), optimizer=opt, net=net, iterator=iterator)
manager = tf.train.CheckpointManager(ckpt, './tf_ckpts', max_to_keep=3)

train_and_checkpoint(net, manager)

tf.train.CheckpointManager 对象会删除旧的检查点。在上面的配置中，它仅保留最近的三个检查点。 

print(manager.checkpoints)  # List the three remaining checkpoints

这些路径（例如 './tf_ckpts/ckpt-10'）并不是磁盘上的文件。它们是 index 文件和一个或多个包含变量值的数据文件的前缀。这些前缀被组合在一个单一的 checkpoint 文件（'./tf_ckpts/checkpoint'）中，CheckpointManager 在该文件中保存其状态。 

ls ./tf_ckpts

加载机制

TensorFlow 通过遍历具有命名边的有向图（从正在加载的对象开始）将变量与检查点值进行匹配。边名称通常来自对象中的属性名称，例如 self.l1 = tf.keras.layers.Dense(5) 中的 "l1"。tf.train.Checkpoint 使用其关键字参数名称，如 tf.train.Checkpoint(step=...) 中的 "step"。

上例中的依赖关系图如下所示

Visualization of the dependency graph for the example training loop

优化器显示为红色，常规变量显示为蓝色，优化器槽变量（slot variables）显示为橙色。其他节点（例如代表 tf.train.Checkpoint 的节点）显示为黑色。

槽变量是优化器状态的一部分，但它们是为特定变量创建的。例如，上图中的 'm' 边对应于动量（momentum），Adam 优化器会为每个变量跟踪动量。只有在变量和优化器都被保存时，槽变量才会保存在检查点中，因此显示为虚线。

在 tf.train.Checkpoint 对象上调用 restore 会将请求的恢复加入队列，一旦从 Checkpoint 对象找到匹配路径，就会立即恢复变量值。例如，您可以通过重构一条穿过网络和层的路径，仅加载您在上面定义的模型中的偏置项（bias）。

to_restore = tf.Variable(tf.zeros([5]))
print(to_restore.numpy())  # All zeros
fake_layer = tf.train.Checkpoint(bias=to_restore)
fake_net = tf.train.Checkpoint(l1=fake_layer)
new_root = tf.train.Checkpoint(net=fake_net)
status = new_root.restore(tf.train.latest_checkpoint('./tf_ckpts/'))
print(to_restore.numpy())  # This gets the restored value.

这些新对象的依赖关系图是您在上面写入的较大检查点的一个非常小的子图。它仅包含偏置项和 tf.train.Checkpoint 用于为检查点编号的保存计数器。

Visualization of a subgraph for the bias variable

restore 返回一个状态对象，该对象具有可选的断言。新 Checkpoint 中创建的所有对象都已恢复，因此 status.assert_existing_objects_matched 通过。 

status.assert_existing_objects_matched()

检查点中有许多未匹配的对象，包括该层的内核（kernel）和优化器的变量。status.assert_consumed 仅在检查点和程序完全匹配时才会通过，在这里会抛出异常。

延迟恢复

TensorFlow 中的 Layer 对象可能会将变量的创建推迟到第一次调用时（即输入形状可用时）。例如，Dense 层内核的形状取决于层的输入和输出形状，因此作为构造函数参数所需的输出形状本身不足以创建变量。由于调用 Layer 也会读取变量的值，因此恢复必须发生在变量创建与其第一次使用之间。

为了支持这种用法，tf.train.Checkpoint 会推迟那些尚无匹配变量的恢复。 

deferred_restore = tf.Variable(tf.zeros([1, 5]))
print(deferred_restore.numpy())  # Not restored; still zeros
fake_layer.kernel = deferred_restore
print(deferred_restore.numpy())  # Restored

手动检查检查点

tf.train.load_checkpoint 返回一个 CheckpointReader，它提供了对检查点内容的更底层访问。它包含从每个变量的键到检查点中每个变量的形状和数据类型的映射。变量的键是其对象路径，如上图所示。 

reader = tf.train.load_checkpoint('./tf_ckpts/')
shape_from_key = reader.get_variable_to_shape_map()
dtype_from_key = reader.get_variable_to_dtype_map()

sorted(shape_from_key.keys())

因此，如果您对 net.l1.kernel 的值感兴趣，可以使用以下代码获取该值

key = 'net/l1/kernel/.ATTRIBUTES/VARIABLE_VALUE'

print("Shape:", shape_from_key[key])
print("Dtype:", dtype_from_key[key].name)

它还提供了一个 get_tensor 方法，允许您检查变量的值。 

reader.get_tensor(key)

对象跟踪

检查点通过“跟踪”其属性中设置的任何变量或可跟踪对象，来保存和恢复 tf.Variable 对象的值。执行保存时，变量会从所有可到达的跟踪对象中递归收集。

与 self.l1 = tf.keras.layers.Dense(5) 等直接属性赋值一样，将列表和字典分配给属性也会跟踪其内容。 

save = tf.train.Checkpoint()
save.listed = [tf.Variable(1.)]
save.listed.append(tf.Variable(2.))
save.mapped = {'one': save.listed[0]}
save.mapped['two'] = save.listed[1]
save_path = save.save('./tf_list_example')

restore = tf.train.Checkpoint()
v2 = tf.Variable(0.)
assert 0. == v2.numpy()  # Not restored yet
restore.mapped = {'two': v2}
restore.restore(save_path)
assert 2. == v2.numpy()

您可能会注意到列表和字典的包装对象。这些包装器是底层数据结构的可检查点版本。就像基于属性的加载一样，这些包装器在添加到容器后会立即恢复变量的值。 

restore.listed = []
print(restore.listed)  # ListWrapper([])
v1 = tf.Variable(0.)
restore.listed.append(v1)  # Restores v1, from restore() in the previous cell
assert 1. == v1.numpy()

可跟踪对象包括 tf.train.Checkpoint、tf.Module 及其子类（例如 keras.layers.Layer 和 keras.Model）以及可识别的 Python 容器

  • dict（以及 collections.OrderedDict）
  • list
  • tuple（以及 collections.namedtuple、typing.NamedTuple）

其他容器类型不支持，包括

  • collections.defaultdict
  • set

所有其他 Python 对象均被忽略，包括

  • int
  • string
  • float

总结

TensorFlow 对象提供了一种简单的自动机制，用于保存和恢复它们所使用变量的值。