Tensorflow 模型持久化

當咱們使用 tensorflow 訓練神經網絡的時候，模型持久化對於咱們的訓練有很重要的做用。

• 若是咱們的神經網絡比較複雜，訓練數據比較多，那麼咱們的模型訓練就會耗時很長，若是在訓練過程當中出現某些不可預計的錯誤，致使咱們的訓練意外終止，那麼咱們將會前功盡棄。爲了不這個問題，咱們就能夠經過模型持久化（保存爲CKPT格式）來暫存咱們訓練過程當中的臨時數據。

• 若是咱們訓練的模型須要提供給用戶作離線的預測，那麼咱們只須要前向傳播的過程，只需獲得預測值就能夠了，這個時候咱們就能夠經過模型持久化（保存爲PB格式）只保存前向傳播中須要的變量並將變量的值固定下來，這個時候只需用戶提供一個輸入，咱們就能夠經過模型獲得一個輸出給用戶。

保存爲 CKPT 格式的模型

1. 定義運算過程
2. 聲明並獲得一個 Saver
3. 經過 Saver.save 保存模型

# coding=UTF-8 支持中文編碼格式
import tensorflow as tf
import shutil
import os.path

MODEL_DIR = "model/ckpt"
MODEL_NAME = "model.ckpt"

# if os.path.exists(MODEL_DIR): 刪除目錄
#     shutil.rmtree(MODEL_DIR)
if not tf.gfile.Exists(MODEL_DIR): #建立目錄
    tf.gfile.MakeDirs(MODEL_DIR)

#下面的過程你能夠替換成CNN、RNN等你想作的訓練過程，這裏只是簡單的一個計算公式
input_holder = tf.placeholder(tf.float32, shape=[1], name="input_holder") #輸入佔位符，並指定名字，後續模型讀取可能會用的
W1 = tf.Variable(tf.constant(5.0, shape=[1]), name="W1")
B1 = tf.Variable(tf.constant(1.0, shape=[1]), name="B1")
_y = (input_holder * W1) + B1
predictions = tf.greater(_y, 50, name="predictions") #輸出節點名字，後續模型讀取會用到，比50大返回true，不然返回false

init = tf.global_variables_initializer()
saver = tf.train.Saver() #聲明saver用於保存模型

with tf.Session() as sess:
    sess.run(init)
    print "predictions : ", sess.run(predictions, feed_dict={input_holder: [10.0]}) #輸入一個數據測試一下
    saver.save(sess, os.path.join(MODEL_DIR, MODEL_NAME)) #模型保存
    print("%d ops in the final graph." % len(tf.get_default_graph().as_graph_def().node)) #獲得當前圖有幾個操做節點

for op in tf.get_default_graph().get_operations(): #打印模型節點信息
    print (op.name, op.values())

運行後生成的文件以下：

• checkpoint ： 記錄目錄下全部模型文件列表
• ckpt.data ： 保存模型中每一個變量的取值
• ckpt.meta ： 保存整個計算圖的結構

保存爲 PB 格式模型

1. 定義運算過程
2. 經過 get_default_graph().as_graph_def() 獲得當前圖的計算節點信息
3. 經過 graph_util.convert_variables_to_constants 將相關節點的values固定
4. 經過 tf.gfile.GFile 進行模型持久化

# coding=UTF-8
import tensorflow as tf
import shutil
import os.path
from tensorflow.python.framework import graph_util


# MODEL_DIR = "model/pb"
# MODEL_NAME = "addmodel.pb"

# if os.path.exists(MODEL_DIR): 刪除目錄
#     shutil.rmtree(MODEL_DIR)
#
# if not tf.gfile.Exists(MODEL_DIR): #建立目錄
#     tf.gfile.MakeDirs(MODEL_DIR)

output_graph = "model/pb/add_model.pb"

#下面的過程你能夠替換成CNN、RNN等你想作的訓練過程，這裏只是簡單的一個計算公式
input_holder = tf.placeholder(tf.float32, shape=[1], name="input_holder")
W1 = tf.Variable(tf.constant(5.0, shape=[1]), name="W1")
B1 = tf.Variable(tf.constant(1.0, shape=[1]), name="B1")
_y = (input_holder * W1) + B1
# predictions = tf.greater(_y, 50, name="predictions") #比50大返回true，不然返回false
predictions = tf.add(_y, 10,name="predictions") #作一個加法運算

init = tf.global_variables_initializer()

with tf.Session() as sess:
    sess.run(init)
    print "predictions : ", sess.run(predictions, feed_dict={input_holder: [10.0]})
    graph_def = tf.get_default_graph().as_graph_def() #獲得當前的圖的 GraphDef 部分，經過這個部分就能夠完成重輸入層到輸出層的計算過程

    output_graph_def = graph_util.convert_variables_to_constants(  # 模型持久化，將變量值固定
        sess,
        graph_def,
        ["predictions"] #須要保存節點的名字
    )
    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))
    print (predictions)

# for op in tf.get_default_graph().get_operations(): 打印模型節點信息
#     print (op.name)

*GraphDef：這個屬性記錄了tensorflow計算圖上節點的信息。

• add_model.pb ： 裏面保存了重輸入層到輸出層這個計算過程的計算圖和相關變量的值，咱們獲得這個模型後傳入一個輸入，既能夠獲得一個預估的輸出值

CKPT 轉換成 PB格式

1. 經過傳入 CKPT 模型的路徑獲得模型的圖和變量數據
2. 經過 import_meta_graph 導入模型中的圖
3. 經過 saver.restore 從模型中恢復圖中各個變量的數據
4. 經過 graph_util.convert_variables_to_constants 將模型持久化

# coding=UTF-8
import tensorflow as tf
import os.path
import argparse
from tensorflow.python.framework import graph_util

MODEL_DIR = "model/pb"
MODEL_NAME = "frozen_model.pb"

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) #檢查目錄下ckpt文件狀態是否可用
    input_checkpoint = checkpoint.model_checkpoint_path #得ckpt文件路徑
    output_graph = os.path.join(MODEL_DIR, MODEL_NAME) #PB模型保存路徑

    output_node_names = "predictions" #原模型輸出操做節點的名字
    saver = tf.train.import_meta_graph(input_checkpoint + '.meta', clear_devices=True) #獲得圖、clear_devices ：Whether or not to clear the device field for an `Operation` or `Tensor` during import.

    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]}) # 測試讀出來的模型是否正確，注意這裏傳入的是輸出 和輸入 節點的 tensor的名字，不是操做節點的名字

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

部分參考：

TensorFlow實戰Google深度學習框架、http://blog.csdn.net/lujiandong1/article/details/53385092