TensorFlow 模型持久化

这里所提到的持久化，是简单的将模型中的计算图进行保存，方便下次直接使用模型，代码主要来源于《Tensorflow 实战Google深度学习框架》。

1、.ckpt文件保存方法

TensorFlow提供的tf.train.Saver类可以很简单的实现模型的保存和还原，一个简单的加法计算图存储的方法如下：

saver.py

import tensorflow as tf
# 变量声明和计算定义
v1 = tf.Variable(tf.constant(1.0, shape = [1]), name = "v1")
v2 = tf.Variable(tf.constant(2.0, shape = [1]), name = "v2")
result = v1 + v2
# 初始化操作定义和持久化模型的Saver类定义
init_op = tf.global_variables_initializer()
saver = tf.train.Saver()
# 结果和计算图的保存
with tf.Session() as sess:
    sess.run(init_op)
    print(sess.run(result))
    saver.svae(sess, "demo_model/model.ckpt")

持久化的优点在于方便随时可以重新使用计算图，只要加载进来就好，加载的方法如下：

load.py

saver = tf.train.import_meta_graph("demo_model/model.ckpt.meta")
with tf.Session() as sess:
    saver.restore(sess, "demo_model/model.ckpt")
    print (sess.run(result))
    print (sess.run(tf.get_default_graph().get_tensor_by_name("add:0")) )

2、滑动平均类的保存

import tensorflow as tf

v = tf.Variable(0, dtype=tf.float32, name="v")
for variables in tf.global_variables(): 
	print (variables.name)
    
ema = tf.train.ExponentialMovingAverage(0.99)
maintain_averages_op = ema.apply(tf.global_variables())
for variables in tf.global_variables(): 
	print (variables.name)

#保存滑动平均模型 
saver = tf.train.Saver()
with tf.Session() as sess:
    init_op = tf.global_variables_initializer()
    sess.run(init_op)
    
    sess.run(tf.assign(v, 10))
    sess.run(maintain_averages_op)
    # 保存的时候会将v:0  v/ExponentialMovingAverage:0这两个变量都存下来。
    saver.save(sess, "Saved_model/model2.ckpt")
    print sess.run([v, ema.average(v)]

'''
可以加载上述数据的方法
v = tf.Variable(0, dtype=tf.float32, name="v")

# 通过变量重命名将原来变量v的滑动平均值直接赋值给v。
# 用法和feed_dict相同
saver = tf.train.Saver({"v/ExponentialMovingAverage": v})
with tf.Session() as sess:
    saver.restore(sess, "Saved_model/model2.ckpt")
    print sess.run(v)
'''

3、.pb文件

tf.train.Saver类需要存储计算图中Tensorflow所有的数据，对于在inference阶段时，只需要将前向传播的输入层到输出层的计算纪录即可。Tensorflow通过提供convert_variables_to_constants函数，可以将变量及其值按照常量的形式存储，做到计算图和变量的合并，存放在同一个文件当中。

saved.py

import tensorflow as tf
from tensorflow.python.framework import graph_util

v1 = tf.Variable(tf.constant(1.0, shape=[1]), name = "v1")
v2 = tf.Variable(tf.constant(2.0, shape=[1]), name = "v2")
result = v1 + v2

init_op = tf.global_variables_initializer()
with tf.Session() as sess:
    sess.run(init_op)
    # 导出计算图的graph_def部分，包含了输入层到输出层的所有计算
    graph_def = tf.get_default_graph().as_graph_def()
    output_graph_def = graph_util.convert_variables_to_constants(sess, graph_def, ['add'])
    # 将导出的模型存入文件
    with tf.gfile.GFile("Saved_model/combined_model.pb", "wb") as f:
           f.write(output_graph_def.SerializeToString())

load.py

from tensorflow.python.platform import gfile
with tf.Session() as sess:
    model_filename = "Saved_model/combined_model.pb"
   # 读取保存的模型文件，并将文件解析成对应的GraphDef Protocol Buffer
    with gfile.FastGFile(model_filename, 'rb') as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())

	# 将graph_def中保存的图加载到当前的图中，return_elements = ["add:0"]给出了返回的张量的名称。
	# 在保存的时候给出的计算节点的名称，所以为“add”。在加载的时候给出的时张量的名称，所以是“add:0”
    result = tf.import_graph_def(graph_def, return_elements=["add:0"])
    print (sess.run(result))