tensorflow 模型与数据的存储与恢复

我们在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，则只输出所有的额变量名，不输出其取值。