TensorFlow 2.0 常用模块1：Checkpoint
文 / 李锡涵，Google Developers Expert
本文节选自《简单粗暴 TensorFlow 2.0》
在上一篇文章中，我们介绍了 TensorFlow 2.0 模型：Keras 训练流程及自定义组件 。
从本篇文章起，我们将开始介绍 TensorFlow 中的常用模块，包括 Checkpoint、TensorBoard、tf.data、@tf.function 等。首先介绍的是 Checkpoint ，让我们来一同了解如何在 TensorFlow 2.0 中保存和恢复模型变量。
tf.train.Checkpoint ：变量的保存与恢复
警告
Checkpoint 只保存模型的参数，不保存模型的计算过程，因此一般用于在具有模型源代码的时候恢复之前训练好的模型参数。如果需要导出模型（无需源代码也能运行模型），请参考 “部署” 章节中的 SavedModel (https://tf.wiki/zh/deployment/export.html#savedmodel)。
很多时候，我们希望在模型训练完成后能将训练好的参数（变量）保存起来。在需要使用模型的其他地方载入模型和参数，就能直接得到训练好的模型。可能你第一个想到的是用 Python 的序列化模块 pickle 存储 model.variables。但不幸的是，TensorFlow 的变量类型 ResourceVariable 并不能被序列化。
好在 TensorFlow 提供了 tf.train.Checkpoint 这一强大的变量保存与恢复类，可以使用其 save() 和 restore() 方法将 TensorFlow 中所有包含 Checkpointable State 的对象进行保存和恢复。具体而言，tf.keras.optimizer 、 tf.Variable 、 tf.keras.Layer 或者 tf.keras.Model 实例都可以被保存。其使用方法非常简单，我们首先声明一个 Checkpoint：
1checkpoint = tf.train.Checkpoint(model=model)
这里 tf.train.Checkpoint() 接受的初始化参数比较特殊，是一个 **kwargs 。具体而言，是一系列的键值对，键名可以随意取，值为需要保存的对象。例如，如果我们希望保存一个继承 tf.keras.Model 的模型实例 model 和一个继承 tf.train.Optimizer 的优化器 optimizer ，我们可以这样写：
1checkpoint = tf.train.Checkpoint(myAwesomeModel=model, myAwesomeOptimizer=optimizer)
这里 myAwesomeModel 是我们为待保存的模型 model 所取的任意键名。注意，在恢复变量的时候，我们还将使用这一键名。
接下来，当模型训练完成需要保存的时候，使用：
1checkpoint.save(save_path_with_prefix)
就可以。 save_path_with_prefix 是保存文件的目录 + 前缀。
注解
例如，在源代码目录建立一个名为 save 的文件夹并调用一次checkpoint.save('./save/model.ckpt')，我们就可以在可以在 save 目录下发现名为checkpoint、model.ckpt-1.index、model.ckpt-1.data-00000-of-00001的三个文件，这些文件就记录了变量信息。checkpoint.save()方法可以运行多次，每运行一次都会得到一个.index 文件和.data 文件，序号依次累加。
当在其他地方需要为模型重新载入之前保存的参数时，需要再次实例化一个 checkpoint，同时保持键名的一致。再调用 checkpoint 的 restore 方法。就像下面这样：
1model_to_be_restored = MyModel() # 待恢复参数的同一模型
2checkpoint = tf.train.Checkpoint(myAwesomeModel=model_to_be_restored) # 键名保持为“myAwesomeModel”
3checkpoint.restore(save_path_with_prefix_and_index)
即可恢复模型变量。 save_path_with_prefix_and_index 是之前保存的文件的目录 + 前缀 + 编号。例如，调用 checkpoint.restore('./save/model.ckpt-1') 就可以载入前缀为 model.ckpt ，序号为 1 的文件来恢复模型。
当保存了多个文件时，我们往往想载入最近的一个。可以使用 tf.train.latest_checkpoint(save_path) 这个辅助函数返回目录下最近一次 checkpoint 的文件名。例如如果 save 目录下有 model.ckpt-1.index 到 model.ckpt-10.index 的 10 个保存文件， tf.train.latest_checkpoint('./save') 即返回 ./save/model.ckpt-10 。
总体而言，恢复与保存变量的典型代码框架如下：
1# train.py 模型训练阶段
2
3model = MyModel()
4# 实例化Checkpoint，指定保存对象为model（如果需要保存Optimizer的参数也可加入）
5checkpoint = tf.train.Checkpoint(myModel=model)
6# ...（模型训练代码）
7# 模型训练完毕后将参数保存到文件（也可以在模型训练过程中每隔一段时间就保存一次）
8checkpoint.save('./save/model.ckpt')
1# test.py 模型使用阶段
2
3model = MyModel()
4checkpoint = tf.train.Checkpoint(myModel=model) # 实例化Checkpoint，指定恢复对象为model
5checkpoint.restore(tf.train.latest_checkpoint('./save')) # 从文件恢复模型参数
6# 模型使用代码
注解
tf.train.Checkpoint与以前版本常用的tf.train.Saver相比，强大之处在于其支持在 Eager Execution 下 “延迟” 恢复变量。具体而言，当调用了checkpoint.restore()，但模型中的变量还没有被建立的时候，Checkpoint 可以等到变量被建立的时候再进行数值的恢复。Eager Execution 下，模型中各个层的初始化和变量的建立是在模型第一次被调用的时候才进行的（好处在于可以根据输入的张量形状而自动确定变量形状，无需手动指定）。这意味着当模型刚刚被实例化的时候，其实里面还一个变量都没有，这时候使用以往的方式去恢复变量数值是一定会报错的。比如，你可以试试在 train.py 调用tf.keras.Model的save_weight()方法保存 model 的参数，并在 test.py 中实例化 model 后立即调用load_weight()方法，就会出错，只有当调用了一遍 model 之后再运行load_weight()方法才能得到正确的结果。可见，tf.train.Checkpoint在这种情况下可以给我们带来相当大的便利。另外，tf.train.Checkpoint同时也支持 Graph Execution 模式。
最后提供一个实例，以前章的 多层感知机模型 为例展示模型变量的保存和载入：
1import tensorflow as tf
2import numpy as np
3import argparse
4from zh.model.mnist.mlp import MLP
5from zh.model.utils import MNISTLoader
6
7parser = argparse.ArgumentParser(description='Process some integers.')
8parser.add_argument('--mode', default='train', help='train or test')
9parser.add_argument('--num_epochs', default=1)
10parser.add_argument('--batch_size', default=50)
11parser.add_argument('--learning_rate', default=0.001)
12args = parser.parse_args()
13data_loader = MNISTLoader()
14
15
16def train():
17 model = MLP()
18 optimizer = tf.keras.optimizers.Adam(learning_rate=args.learning_rate)
19 num_batches = int(data_loader.num_train_data // args.batch_size * args.num_epochs)
20 checkpoint = tf.train.Checkpoint(myAwesomeModel=model) # 实例化Checkpoint，设置保存对象为model
21 for batch_index in range(1, num_batches+1):
22 X, y = data_loader.get_batch(args.batch_size)
23 with tf.GradientTape() as tape:
24 y_pred = model(X)
25 loss = tf.keras.losses.sparse_categorical_crossentropy(y_true=y, y_pred=y_pred)
26 loss = tf.reduce_mean(loss)
27 print("batch %d: loss %f" % (batch_index, loss.numpy()))
28 grads = tape.gradient(loss, model.variables)
29 optimizer.apply_gradients(grads_and_vars=zip(grads, model.variables))
30 if batch_index % 100 == 0: # 每隔100个Batch保存一次
31 path = checkpoint.save('./save/model.ckpt') # 保存模型参数到文件
32 print("model saved to %s" % path)
33
34
35def test():
36 model_to_be_restored = MLP()
37 # 实例化Checkpoint，设置恢复对象为新建立的模型model_to_be_restored
38 checkpoint = tf.train.Checkpoint(myAwesomeModel=model_to_be_restored)
39 checkpoint.restore(tf.train.latest_checkpoint('./save')) # 从文件恢复模型参数
40 y_pred = np.argmax(model_to_be_restored.predict(data_loader.test_data), axis=-1)
41 print("test accuracy: %f" % (sum(y_pred == data_loader.test_label) / data_loader.num_test_data))
42
43
44if __name__ == '__main__':
45 if args.mode == 'train':
46 train()
47 if args.mode == 'test':
48 test()
在代码目录下建立 save 文件夹并运行代码进行训练后，save 文件夹内将会存放每隔 100 个 batch 保存一次的模型变量数据。在命令行参数中加入 --mode=test 并再次运行代码，将直接使用最后一次保存的变量值恢复模型并在测试集上测试模型性能，可以直接获得 95% 左右的准确率。
使用
tf.train.CheckpointManager删除旧的 Checkpoint 以及自定义文件编号
在模型的训练过程中，我们往往每隔一定步数保存一个 Checkpoint 并进行编号。不过很多时候我们会有这样的需求：
在长时间的训练后，程序会保存大量的 Checkpoint，但我们只想保留最后的几个 Checkpoint；
Checkpoint 默认从 1 开始编号，每次累加 1，但我们可能希望使用别的编号方式（例如使用当前 Batch 的编号作为文件编号）。
这时，我们可以使用 TensorFlow 的
tf.train.CheckpointManager来实现以上需求。具体而言，在定义 Checkpoint 后接着定义一个 CheckpointManager：
1checkpoint = tf.train.Checkpoint(model=model)
2manager = tf.train.CheckpointManager(checkpoint, directory='./save', checkpoint_name='model.ckpt', max_to_keep=k)
此处，
directory参数为文件保存的路径，checkpoint_name为文件名前缀（不提供则默认为ckpt），max_to_keep为保留的 Checkpoint 数目。在需要保存模型的时候，我们直接使用
manager.save()即可。如果我们希望自行指定保存的 Checkpoint 的编号，则可以在保存时加入checkpoint_number参数。例如manager.save(checkpoint_number=100)。
以下提供一个实例，展示使用 CheckpointManager 限制仅保留最后三个 Checkpoint 文件，并使用 batch 的编号作为 Checkpoint 的文件编号：
1import tensorflow as tf
2import numpy as np
3import argparse
4from zh.model.mnist.mlp import MLP
5from zh.model.utils import MNISTLoader
6
7parser = argparse.ArgumentParser(description='Process some integers.')
8parser.add_argument('--mode', default='train', help='train or test')
9parser.add_argument('--num_epochs', default=1)
10parser.add_argument('--batch_size', default=50)
11parser.add_argument('--learning_rate', default=0.001)
12args = parser.parse_args()
13data_loader = MNISTLoader()
14
15
16def train():
17 model = MLP()
18 optimizer = tf.keras.optimizers.Adam(learning_rate=args.learning_rate)
19 num_batches = int(data_loader.num_train_data // args.batch_size * args.num_epochs)
20 checkpoint = tf.train.Checkpoint(myAwesomeModel=model)
21 # 使用tf.train.CheckpointManager管理Checkpoint
22 manager = tf.train.CheckpointManager(checkpoint, directory='./save', max_to_keep=3)
23 for batch_index in range(1, num_batches):
24 X, y = data_loader.get_batch(args.batch_size)
25 with tf.GradientTape() as tape:
26 y_pred = model(X)
27 loss = tf.keras.losses.sparse_categorical_crossentropy(y_true=y, y_pred=y_pred)
28 loss = tf.reduce_mean(loss)
29 print("batch %d: loss %f" % (batch_index, loss.numpy()))
30 grads = tape.gradient(loss, model.variables)
31 optimizer.apply_gradients(grads_and_vars=zip(grads, model.variables))
32 if batch_index % 100 == 0:
33 # 使用CheckpointManager保存模型参数到文件并自定义编号
34 path = manager.save(checkpoint_number=batch_index)
35 print("model saved to %s" % path)
36
37
38def test():
39 model_to_be_restored = MLP()
40 checkpoint = tf.train.Checkpoint(myAwesomeModel=model_to_be_restored)
41 checkpoint.restore(tf.train.latest_checkpoint('./save'))
42 y_pred = np.argmax(model_to_be_restored.predict(data_loader.test_data), axis=-1)
43 print("test accuracy: %f" % (sum(y_pred == data_loader.test_label) / data_loader.num_test_data))
44
45
46if __name__ == '__main__':
47 if args.mode == 'train':
48 train()
49 if args.mode == 'test':
50 test()福利 | 问答环节
我们知道在入门一项新的技术时有许多挑战与困难需要克服。如果您有关于 TensorFlow 的相关问题，可在本文后留言，我们的工程师和 GDE 将挑选其中具有代表性的问题在下一期进行回答～
在上一篇文章《TensorFlow 2.0 模型：Keras 训练流程及自定义组件》中，我们对于部分具有代表性的问题回答如下：
Q1. tf.data 如何读取本地大的数据集（几百 G），能出一期教程吗？谢谢！
A：我们会在后面的连载系列中介绍 tf.data 的基本使用方式。针对大规模数据集，可以考虑先处理为 TFRecord 格式，可参考示例展示的如何将图片数据处理为 TFRecord 文件并读取。
https://tensorflow.google.cn/tutorials/load_data/tfrecord#walkthrough_reading_and_writing_image_data
Q2. linux 的包还没 ready 吗，为什么 pip 默认安装的还是 1.14？
A：请检查您所使用的 pip 源，可能是您所用的 pip 的源尚未更新。
Q3. 能不能写一篇控制网络参数更新或者控制某一层内部分参数进行梯度更新其他停止更新的帖子。不是很懂？
A：在 TensorFlow 2.0 中，可以非常方便地实现部分参数更新。您在调用tape.gradient时可以仅指定那些需要进行梯度更新的变量。示例如下： 1import tensorflow as tf
2from zh.model.mnist.mlp import MLP
3from zh.model.utils import MNISTLoader
4
5num_epochs = 5
6batch_size = 50
7learning_rate_1 = 0.001
8learning_rate_2 = 0.01
9
10model = MLP()
11data_loader = MNISTLoader()
12# 声明两个优化器，设定不同的学习率，分别用于更新MLP模型的第一层和第二层
13optimizer_1 = tf.keras.optimizers.Adam(learning_rate=learning_rate_1)
14optimizer_2 = tf.keras.optimizers.Adam(learning_rate=learning_rate_2)
15num_batches = int(data_loader.num_train_data // batch_size * num_epochs)
16for batch_index in range(num_batches):
17 X, y = data_loader.get_batch(batch_size)
18 with tf.GradientTape(persistent=True) as tape: # 声明一个持久的GradientTape，允许我们多次调用tape.gradient方法
19 y_pred = model(X)
20 loss = tf.keras.losses.sparse_categorical_crossentropy(y_true=y, y_pred=y_pred)
21 loss = tf.reduce_mean(loss)
22 print("batch %d: loss %f" % (batch_index, loss.numpy()))
23 grads = tape.gradient(loss, model.dense1.variables) # 单独求第一层参数的梯度
24 optimizer_1.apply_gradients(grads_and_vars=zip(grads, model.dense1.variables)) # 单独对第一层参数更新，学习率0.001
25 grads = tape.gradient(loss, model.dense2.variables) # 单独求第二层参数的梯度
26 optimizer_1.apply_gradients(grads_and_vars=zip(grads, model.dense2.variables)) # 单独对第二层参数更新，学习率0.01Q4. 能够提供一些 a 卡支持 TensorFlow 的信息吗？
A：可参考 使用 AMD 显卡加速 TensorFlow。
Q5. 有没有使用 tfrecords 完整的教程？
A：可以参考
https://tensorflow.google.cn/tutorials/load_data/tfrecord
《简单粗暴 TensorFlow 2.0 》目录
公众号回复关键字“手册”获取系列内容合集及 FAQ。