【Tensorflow教程笔记】常用模块 tf.train.Checkpoint ：变量的保存与恢复

最新推荐文章于 2022-05-21 15:21:31 发布

_APTX4869

最新推荐文章于 2022-05-21 15:21:31 发布

阅读量8.9k

点赞数 12

分类专栏： Tensorflow

本文链接：https://blog.csdn.net/nanke_4869/article/details/114203871

版权

本文详细介绍了 TensorFlow 中如何使用 tf.train.Checkpoint 类来保存和恢复模型参数。通过实例展示了保存、加载模型的步骤，包括如何在即时执行模式下延迟恢复变量，以及如何利用 CheckpointManager 控制 Checkpoint 文件数量和自定义文件编号。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

tf.train.Checkpoint

很多时候，我们希望在模型训练完成后能将训练好的参数（变量）保存起来。在需要使用模型的其他地方载入模型和参数，就能直接得到训练好的模型。可能你第一个想到的是用 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：

checkpoint = tf.train.Checkpoint(model=model)

这里 tf.train.Checkpoint() 接受的初始化参数比较特殊，是一个 **kwargs 。具体而言，是一系列的键值对，键名可以随意取，值为需要保存的对象。例如，如果我们希望保存一个继承 tf.keras.Model 的模型实例 model 和一个继承 tf.train.Optimizer 的优化器 optimizer ，我们可以这样写：

checkpoint = tf.train.Checkpoint(myAwesomeModel=model, myAwesomeOptimizer=optimizer)

这里 myAwesomeModel 是我们为待保存的模型 model 所取的任意键名。注意，在恢复变量的时候，我们还将使用这一键名。

保存参数

接下来，当模型训练完成需要保存的时候，使用：

checkpoint.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 方法。就像下面这样：

model_to_be_restored = MyModel()                                        # 待恢复参数的同一模型
checkpoint = tf.train.Checkpoint(myAwesomeModel=model_to_be_restored)   # 键名保持为“myAwesomeModel”
checkpoint.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 。

保存变量+恢复变量

总体而言，恢复与保存变量的典型代码框架如下：

# train.py 模型训练阶段

model = MyMod

最低0.47元/天解锁文章