我们在tensorflow中训练模型时,往往一个模型需要训练好几天,训练中可能出现意外关机而停止训练。如果此时再重头开始,必然使得前面已经训练的结果功亏一篑。于是,就有了save与restore。变量在文件中以name为名称存储
import tensorflow as tf
# 定义存储路径
save_path = "./test/model.ckpt"
# 存储变量
v1 = tf.get_variable(name="v1", shape=[3], initializer=tf.zeros_initializer)
v2 = tf.get_variable(name="v2", shape=[5], initializer=tf.zeros_initializer)
inc_v1 = v1.assign(v1+1)
dec_v2 = v2.assign(v2-1)
init_op = tf.global_variables_initializer()
saver = tf.train.Saver()
with tf.Session() as sess:
sess.run(init_op)
inc_v1.op.run()
dec_v2.op.run()
save_path = saver.save(sess,save_path=save_path)
上面代码执行以后,将在save_path文件夹下生成四个文件,分别命名为:
ckeckpoint
model.ckpt.data-00000-of-00001
model.ckpt.index
model.ckpt.meta
此时,若全部注释掉上面的代码,再运行下面的代码,将会恢复所有变量
import tensorflow as tf
# 定义存储路径
save_path = "./test/model.ckpt"
tf.reset_default_graph()
v1 = tf.get_variable(name="v1", shape=[3], initializer=tf.zeros_initializer)
v2 = tf.get_variable(name="v2", shape=[5], initializer=tf.zeros_initializer)
saver = tf.train.Saver()
with tf.Session() as sess:
saver.restore(sess,save_path)
print(v1.eval(), v2.eval())
但是,如果我们多定义了一个变量v3,再运行如下代码,将会报错:
import tensorflow as tf
# 定义存储路径
save_path = "./test/model.ckpt"
tf.reset_default_graph()
v1 = tf.get_variable(name="v1", shape=[3], initializer=tf.zeros_initializer)
v2 = tf.get_variable(name="v2", shape=[5], initializer=tf.zeros_initializer)
v3 = tf.get_variable(name="v3", shape=[5], initializer=tf.zeros_initializer)
saver = tf.train.Saver()
with tf.Session() as sess:
v3.initializer.run()
saver.restore(sess,save_path)
print(v1.eval(), v2.eval(),v3.eval())
错误为:NotFoundError (see above for traceback): Key v3 not found in checkpoint.即,名字为v3的变量找不到。也即是说,模型的变量数不能多于checkpoint中存储的变量数。
但是,如果我们只训练了一个五层网络,将这五层网络的变量存储到了checkpoint,现在我们要拓展网络到第六层,想用之前训练的五层网络的参数来初始化现在的第六层网络,怎么办呢?于是就遇到了我们需要存储与恢复部分变量的问题。我们可以以列表或者字典的形式在tf.train.Saver()中传递我们要存储与恢复的参数,接着上面代码中生成的checkpoint文件,继续运行如下代码:
tf.reset_default_graph()
v1 = tf.get_variable("v1", [3], initializer = tf.zeros_initializer)
v2 = tf.get_variable("v2", [5], initializer = tf.zeros_initializer)
# 也可以写作saver = tf.train.Saver([v2])
saver = tf.train.Saver({"v2":v2})
# Use the saver object normally after that.
with tf.Session() as sess:
# 既然不需要从checkpoint文件中恢复v1,那么就需要手动初始化v1
v1.initializer.run()
saver.restore(sess,save_path)
print("v1 : %s" % v1.eval())
print("v2 : %s" % v2.eval())
输出 :
v1 : [0. 0. 0.]
v2 : [-1. -1. -1. -1. -1.]
可见,只有v2是从checkpoint中恢复的。
我们可以定义多个Saver()对象,也可以将同一个变量挂在多个Saver()对象里面。
有时候,我们想知道checkpoint里面存储了哪些变量,或者他们的值是多少,或者只关心给定的变量,可用如下代码:
from tensorflow.python.tools import inspect_checkpoint as chkp
chkp.print_tensors_in_checkpoint_file(save_path, tensor_name='', all_tensors=True,all_tensor_names=True)
print("***************************************")
chkp.print_tensors_in_checkpoint_file(save_path, tensor_name='v1', all_tensors=False,all_tensor_names=False)
print("***************************************")
chkp.print_tensors_in_checkpoint_file(save_path, tensor_name='v2', all_tensors=False,all_tensor_names=True)
输出:
tensor_name: v1
[1. 1. 1.]
tensor_name: v2
[-1. -1. -1. -1. -1.]
***************************************
tensor_name: v1
[1. 1. 1.]
***************************************
tensor_name: v1
tensor_name: v2
如果tensor_name="",当all_tensors=True时输出所有的tensors名及相应的取值。若all_tensors=False及all_tensor_names=True,则只输出所有的额变量名,不输出其取值。