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") #模型目錄