模型的save和restore

保存模型

import tensorflow as tf

v1 = tf.Variable(1.0, name='v1')
v2 = tf.Variable(2.0, name='v2')

result = v1 + v2

saver = tf.train.Saver()

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    print(sess.run(result))    # 3.0
    saver.save(sess,'./model/model.ckpt')

保存的文件如下：

可以直接加载已经持久化的图

import tensorflow as tf

saver = tf.train.import_meta_graph('./model/model.ckpt.meta')

with tf.Session() as sess:
    saver.restore(sess, './model/model.ckpt')
    print(sess.run(tf.get_default_graph().get_tensor_by_name('add:0')))  # 3.0

也可以重新定义图

import tensorflow as tf

v1 = tf.Variable(1.0,'renamed-v1')
v2 = tf.Variable(2.0,'renamed-v2')

result = v1 + v2

# 下面这句不对,因为保存变量的名称和加载变量时的名称不一致,即'v1'和'renamed-v1'不一致
#saver = tf.train.Saver()
# 使用字典对变量重命名就可以加载原来的模型了,该字典指定了原来名称为‘v1’的变量v1加载到名称为'renamed-v1'的变量v1中,v2同理
saver = tf.train.Saver({'v1':v1, 'v2':v2})

with tf.Session() as sess:
    saver.restore(sess,'./model/model.ckpt')
    print(sess.run(result))

加载部分变量

上面的三段代码默认保存和加载了计算图上定义的全部变量，但有时可能只需要保存或加载部分变量。比如，有一个训练好的5层的神经网络模型，现在想尝试一个6层的神经网络，那么可以将前面5层神经网络中的参数直接加载到新的模型，而仅将最后一层神经网络重新训练即可。
为了保存或者加载部分变量，在声明tf.train.Saver类时可以提供一个列表来指定需要保存或者加载的变量。比如在加载模型的代码中使用saver=tf.train.Saver()命令来构建tf.train.Saver类，那么只有v1会被加载进来。如果运行修改后只加载了v1的代码会得到尝试使用未初始化变量的错误

import tensorflow as tf

v1 = tf.Variable(1.0,name='v1')
v2 = tf.Variable(2.0,name='v2')

result = v1 + v2

saver = tf.train.Saver([v1])

with tf.Session() as sess:
    saver.restore(sess,'./model/model.ckpt')
    print(sess.run(v1))
    print(sess.run(result)) # 报错：Attempting to use uninitialized value v2

加载模型时对变量重命名便于使用滑动平均值

将变量重命名的主要目的之一是方便使用变量的滑动平均值,在TensorFlow中每一个变量的滑动平均值是通过影子变量维护的,所以要获取变量的滑动平均值实际上就是获取这个影子变量的取值,如果在加载模型时直接将影子变量映射到变量自身,在使用训练好的模型时就不需要再调用滑动平均函数来计算滑动平均值了,这样大大方便滑动平均模型的使用。

# 保存变量和其滑动平均值
import tensorflow as tf

v = tf.Variable(0.1, name='v')
for var in tf.global_variables():
    print( var.name) # 此处只会输出'v:0'

ema = tf.train.ExponentialMovingAverage(0.999)
avg_op = ema.apply(tf.global_variables()) # TensorFlow自动生成一个影子变量

for var in tf.global_variables():
    print( var.name) # 输出'v:0'和'v/ExponentialMovingAverage:0'

saver = tf.train.Saver()

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    sess.run(tf.assign(v,10))
    sess.run(avg_op)
    saver.save(sess,'./model/model.ckpt') # 'v:0'和'v/ExponentialMovingAverage:0'两个变量都会被保存
    print(sess.run([v, ema.average(v)]))  # 输出: [10.0, 0.10989987]

通过变量重命名的方式直接将变量的滑动平均值赋值给变量

import tensorflow as tf

v = tf.Variable(0.1, name='v')
saver = tf.train.Saver({'v/ExponentialMovingAverage':v}) # 字典的含义是：需要restore的变量的name:被赋值的变量
with tf.Session() as sess:
    saver.restore(sess,'./model/model.ckpt') # 'v:0'和'v/ExponentialMovingAverage:0'两个变量都会被保存
    print(sess.run([v]))                     # 输出：[0.10989987]

方便加载模型时重命名滑动平均变量起见，tf.train.ExponentialMovingAverage类提供了variables_to_restore函数来生成tf.train.Saver类所需要的变量重命名字典。

import tensorflow as tf

v = tf.Variable(0.1, name='v')
ema = tf.train.ExponentialMovingAverage(0.999) # 此处的decay取值不影响v的值但最好写成和训练时一致

print(ema.variables_to_restore()) #输出:{'v/ExponentialMovingAverage': <tf.Variable 'v:0' shape=() dtype=float32_ref>}

saver = tf.train.Saver(ema.variables_to_restore())
with tf.Session() as sess:
    saver.restore(sess,'./model/model.ckpt') # 'v:0'和'v/ExponentialMovingAverage:0'两个变量都会被保存
    print(sess.run([v]))  # 输出：[0.10989987]

Reference

1. 郑泽宇等.TensorFLow实战Google深度学习框架(第2版),电子工业出版社,2018.