官方中文文档的网址先贴出来:https://tensorflow.google.cn/programmers_guide/saved_model
tf.train.Saver
类别提供了保存和恢复模型的方法。tf.train.Saver
构造函数针对图中所有变量或指定列表的变量将 save
和 restore
op 添加到图中。Saver
对象提供了运行这些 op 的方法,指定了写入或读取检查点文件的路径。
TensorFlow 将变量保存在二进制检查点文件中,简略而言,这类文件将变量名称映射到张量值。
保存变量
使用 tf.train.Saver()
创建 Saver
来管理模型中的所有变量。例如,以下代码片段展示了如何调用 tf.train.Saver.save
方法以将变量保存到检查点文件中:
# Create some variables.
v1 = tf.get_variable("v1", shape=[3], initializer = tf.zeros_initializer)
v2 = tf.get_variable("v2", shape=[5], initializer = tf.zeros_initializer)
inc_v1 = v1.assign(v1+1)
dec_v2 = v2.assign(v2-1)
# Add an op to initialize the variables.
init_op = tf.global_variables_initializer()
# Add ops to save and restore all the variables.
saver = tf.train.Saver()
# Later, launch the model, initialize the variables, do some work, and save the
# variables to disk.
with tf.Session() as sess:
sess.run(init_op)
# Do some work with the model.
inc_v1.op.run()
dec_v2.op.run()
# Save the variables to disk.
save_path = saver.save(sess, "/tmp/model.ckpt")
print("Model saved in path: %s" % save_path)
恢复变量
tf.train.Saver
对象不仅将变量保存到检查点文件中,还将恢复变量。请注意,当您恢复变量时,您不必事先将其初始化。例如,以下代码片段展示了如何调用 tf.train.Saver.restore
方法以从检查点文件中恢复变量:
tf.reset_default_graph()
# Create some variables.
v1 = tf.get_variable("v1", shape=[3])
v2 = tf.get_variable("v2", shape=[5])
# Add ops to save and restore all the variables.
saver = tf.train.Saver()
# Later, launch the model, use the saver to restore variables from disk, and
# do some work with the model.
with tf.Session() as sess:
# Restore variables from disk.
saver.restore(sess, "/tmp/model.ckpt")
print("Model restored.")
# Check the values of the variables
print("v1 : %s" % v1.eval())
print("v2 : %s" % v2.eval())
选择要保存和恢复的变量(选择网络中的部分变量保存或者恢复)
如果您没有向 tf.train.Saver()
传递任何参数,则 Saver 会处理图中的所有变量。每个变量都保存在创建变量时所传递的名称下。
在检查点文件中明确指定变量名称的这种做法有时会非常有用。例如,您可能已经使用名为"weights"
的变量训练了一个模型,而您想要将该变量的值恢复到名为"params"
的变量中。
有时候,仅保存或恢复模型使用的变量子集也会很有裨益。例如,您可能已经训练了一个五层的神经网络,现在您想要训练一个六层的新模型,并重用该五层的现有权重。您可以使用 Saver 只恢复这前五层的权重。
您可以向 tf.train.Saver()
的构造函数传递以下任一内容来轻松指定要保存或加载的名称和变量:
- 变量列表(将以其本身的名称保存)。
- Python 字典,其中,键是要使用的名称,键值是要管理的变量。
继续前面所示的保存/恢复示例:
tf.reset_default_graph()
# Create some variables.
v1 = tf.get_variable("v1", [3], initializer = tf.zeros_initializer)
v2 = tf.get_variable("v2", [5], initializer = tf.zeros_initializer)
# Add ops to save and restore only `v2` using the name "v2"
saver = tf.train.Saver({"v2": v2})
# Use the saver object normally after that.
with tf.Session() as sess:
# Initialize v1 since the saver will not.
v1.initializer.run()
saver.restore(sess, "/tmp/model.ckpt")
print("v1 : %s" % v1.eval())
print("v2 : %s" % v2.eval())
注意:
-
如果需要保存和恢复模型变量的不同子集,您可以根据需要创建任意数量的
Saver
对象。同一个变量可以列在多个 Saver 对象中,变量的值只有在Saver.restore()
方法运行时才会更改。 -
如果您仅在会话开始时恢复模型变量的子集,则必须为其他变量运行初始化 op。有关详情,请参阅
tf.variables_initializer
。 -
要检查某个检查点的变量,您可以使用
inspect_checkpoint
库,尤其是print_tensors_in_checkpoint_file
函数。 -
默认情况下,
Saver
会为每个变量使用tf.Variable.name
属性的值。但是,当您创建一个Saver
对象时,您可以选择为检查点文件中的变量选择名称(此为可选操作)。
检查某个检查点的变量
可以使用 inspect_checkpoint
库快速检查某个检查点的变量。
# import the inspect_checkpoint library
from tensorflow.python.tools import inspect_checkpoint as chkp
# print all tensors in checkpoint file
chkp.print_tensors_in_checkpoint_file("/tmp/model.ckpt", tensor_name='', all_tensors=True)
# tensor_name: v1
# [ 1. 1. 1.]
# tensor_name: v2
# [-1. -1. -1. -1. -1.]
# print only tensor v1 in checkpoint file
chkp.print_tensors_in_checkpoint_file("/tmp/model.ckpt", tensor_name='v1', all_tensors=False)
# tensor_name: v1
# [ 1. 1. 1.]
# print only tensor v2 in checkpoint file
chkp.print_tensors_in_checkpoint_file("/tmp/model.ckpt", tensor_name='v2', all_tensors=False)
# tensor_name: v2
# [-1. -1. -1. -1. -1.]
保存和恢复模型概述
如果您想保存和加载变量、图,以及图的元数据 - 简而言之,如果您想保存或恢复模型 - 我们推荐使用 SavedModel。SavedModel 是一种与语言无关,可恢复的密封式序列化格式。SavedModel 可让较高级别的系统和工具创建、使用和变换 TensorFlow 模型。TensorFlow 提供了多种与 SavedModel 交互的机制,如 tf.saved_model API、Estimator API 和 CLI。
用于构建和加载 SavedModel 的 API