深度学习：Tensorflow模型持久化

Tensorflow训练的模型，常常需要保存下来，便于调用，有两种方式保存：ckpt和pb

1、保存为ckpt模型，代码为：

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

init_op = tf.global_variables_initializer()
saver = tf.train.Saver()

with tf.Session() as sess:
    sess.run(init_op)
    saver.save(sess,"./model.ckpt")

保存的文件夹中，有4个文件：

checkpoint:保存了一个目录下所有的模型文件列表

model.ckpt.data-00000-of-00001:

model.ckpt.index:保存了Tensorflow程序中每一个变量的取值

model.ckpt.meta:保存Tensorflow计算图的结构

2、调用ckpt模型，代码为：

import tensorflow as tf

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

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

3、将计算图结构和变量参数值统一存放在一个文件中，保存为pb模型，代码为：

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()
saver = tf.train.Saver()

with tf.Session() as sess:
    sess.run(init_op)

    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("./model.pb","wb") as f:
        f.write(output_graph_def.SerializeToString())
    

4、调用pb模型，代码为：

import tensorflow as tf
from tensorflow.python.platform import gfile

with tf.Session() as sess:
    model_filename = "model.pb"
    with gfile.FastGFile(model_filename,'rb') as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())

    result = tf.import_graph_def(graph_def,return_elements = ["add:0"])
    print(sess.run(result))