Learning Tensorflow(3)--- 模型持久化

通过模型持久化（保存为CKPT格式）来暂存我们训练过程中的临时数据。

 

通过模型持久化（保存为PB格式）只保存前向传播中需要的变量并将变量的值固定下来，这个时候只需用户提供一个输入，我们就可以通过模型得到一个输出给用户。

CKTP格式

保存模型

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.initialze_all_variables()
Saver = tf.train.Saver()
With tf.Session() as sess:
         Sess.run(init_op)
         Saver.save(sess, “/path/to/model/model.ckpt”)

 

上面这段代码生成三个文件，分别为model.ckpt.meta, model.ckpt,checkpoint.

• model.ckpt.meta保存了TensorFlow计算图的结构.
• Model.ckpt保存了TensorFlow程序中每一个变量的取值。
• Checkpoint保存了一个目录下所有的模型文件列表。

 

还原模型

定义计算图后，还原变量值

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 = tf.train.Saver()
With tf.Session() as sess:
         Saver.save(sess, “/path/to/model/model.ckpt”)
         Print sess.run(result)

加载持久化模型

Import tensorflow as tf
Saver = tf.train.import_meta_graph("model.ckpt.meta")
With tf.Session() as sess:
         Saver.resotre(sess, “model.ckpt-18000”)
         Print sess.run(tf.get_default_graph().get_tensor_by_name(“add:0”))

 

PB格式

保存为PB格式

constant_graph = graph_util.convert_variables_to_constants(sess, sess.graph_def, ["output"])
with tf.gfile.FastGFile(pb_file_path, mode='wb') as f:
f.write(constant_graph.SerializeToString())

 

加载pb模型

完整的加载模型，并利用模型来进行分类

import tensorflow as tf
import  numpy as np
import PIL.Image as Image
from skimage import io, transform
def classification(jpg_path, pb_file_path):
    with tf.Graph().as_default():
        output_graph_def = tf.GraphDef()
        with open(pb_file_path, "rb") as f:
            output_graph_def.ParseFromString(f.read())
            _ = tf.import_graph_def(output_graph_def, name="")
        with tf.Session() as sess:
            init = tf.global_variables_initializer()
            sess.run(init)
            input_x = sess.graph.get_tensor_by_name("input:0")
            print input_x
            out_softmax = sess.graph.get_tensor_by_name("softmax:0")
            print out_softmax
            out_label = sess.graph.get_tensor_by_name("output:0")
            print out_label
            img = io.imread(jpg_path)
            img = transform.resize(img, (224, 224, 3))
            img_out_softmax = sess.run(out_softmax, feed_dict={input_x:np.reshape(img, [-1, 224, 224, 3])})
            print "img_out_softmax:",img_out_softmax
            prediction_labels = np.argmax(img_out_softmax, axis=1)
            print "label:",prediction_labels

CKTP格式转PB格式

通过载入cktp模型得到计算图，通过restore恢复计算图中各个变量的值，然后在通过graph_util.convert_variables_to_constants将模型持久化，写入PB文件。
 

# coding=UTF-8

import tensorflow as tf
import os.path
import argparse
from tensorflow.python.framework import graph_util

if not tf.gfile.Exists(MODEL_DIR): #创建目录
    tf.gfile.MakeDirs(MODEL_DIR)

def freeze_graph(model_folder):
    checkpoint = tf.train.get_checkpoint_state(model_folder)
    input_checkpoint = checkpoint.model_checkpoint_path
    output_graph = os.path.join(MODEL_DIR, MODEL_NAME)
    output_node_names = "predictions"
    saver = tf.train.import_meta_graph(input_checkpoint + '.meta', clear_devices=True)
    graph = tf.get_default_graph()
    input_graph_def = graph.as_graph_def()

    with tf.Session() as sess:
        saver.restore(sess, input_checkpoint) #恢复图并得到数据
        print "predictions : ", sess.run("predictions:0", feed_dict={"input_holder:0": [10.0]})
        output_graph_def = graph_util.convert_variables_to_constants(
            sess,
            input_graph_def,
            output_node_names.split(",")
        )
        with tf.gfile.GFile(output_graph, "wb") as f:
            f.write(output_graph_def.SerializeToString())
        print("%d ops in the final graph." % len(output_graph_def.node)) #得到当前图有几个操作节点
        for op in graph.get_operations():
            print(op.name, op.values())

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("model_folder", type=str, help="input ckpt model dir") #命令行解析，help是提示符，type是输入的类型，

    # 这里运行程序时需要带上模型ckpt的路径，不然会报 error: too few arguments
    aggs = parser.parse_args()
    freeze_graph(aggs.model_folder)
  # freeze_graph("model/ckpt") #模型目录