使用tensorflow訓練模型的時候，模型持久化對我們來說非常重要。
如果我們的模型比較複雜，需要的資料比較多，那麼在模型的訓練時間會耗時很長。如果在訓練過程中出現了模型不可預期的錯誤，導緻訓練意外終止，那麼我們将會前功盡棄。為了解決這一問題，我們可以使用模型持久化(儲存為ckpt檔案格式)來儲存我們在訓練過程中的臨時資料。、
如果我們訓練出的模型需要提供給使用者做離線預測，那麼我們隻需要完成前向傳播過程。這個時候我們就可以使用模型持久化(儲存為pb檔案格式)來隻儲存前向傳播過程中的變量并将變量固定下來，這時候使用者隻需要提供一個輸入即可得到前向傳播的預測結果。
ckpt和pb持久化方式的差別在于ckpt檔案将模型結構與模型權重分離儲存，便于訓練過程;pb檔案則是graph_def的序列化檔案，便于釋出和離線預測。官方提供freeze_grpah.py腳本來将ckpt檔案轉為pb檔案。
CKPT模型持久化
首先定義前向傳播過程;
聲明并得到一個Saver;
使用Saver.save()儲存模型;
# coding=UTF-8 支援中文編碼格式
import tensorflow as tf
import shutil
import os.path
MODEL_DIR = "/home/zheng/PycharmProjects/ckptLoad/Models/"
MODEL_NAME = "model.ckpt"
#下面的過程你可以替換成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.add(_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)) #得到目前圖有幾個操作節點
predictions : [ 101.]
28 ops in the final graph.
注：代碼含義請參考注釋，需要注意的是可以自定義模型儲存的路徑
ckpt模型持久化使用起來非常簡單，隻需要我們聲明一個tf.train.Saver，然後調用save()函數，将會話模型儲存到指定的目錄。執行代碼結果，會在我們指定模型目錄下出現4個檔案
checkpoint ： 記錄目錄下所有模型檔案清單
ckpt.data ： 儲存模型中每個變量的取值
ckpt.meta ： 儲存整個計算圖的結構
ckpt模型加載
# -*- coding: utf-8 -*-)
import tensorflow as tf
from numpy.random import RandomState
# 定義訓練資料batch的大小
batch_size = 8
#下面的過程你可以替換成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.add(_y, 50, name="predictions") #輸出節點名字，後續模型讀取會用到，比50大傳回true，否則傳回false
#saver=tf.train.Saver()
# creare a session，建立一個會話來運作TensorFlow程式
with tf.Session() as sess:
saver = tf.train.import_meta_graph('/home/zheng/Models/model/model.meta')
saver.restore(sess, tf.train.latest_checkpoint('/home/zheng/Models/model'))
#saver.restore(sess, tf.train.latest_checkpoint('/home/zheng/Models/model'))
# 初始化變量
sess.run(tf.global_variables_initializer())
print "predictions : ", sess.run(predictions, feed_dict={input_holder: [10.0]})
代碼結果，可以看到運作結果一樣
predictions : [ 101.]
PB模型持久化
定義運算過程
通過 get_default_graph().as_graph_def() 得到目前圖的計算節點資訊
通過 graph_util.convert_variables_to_constants 将相關節點的values固定
通過 tf.gfile.GFile 進行模型持久化
# coding=UTF-8
import tensorflow as tf
import shutil
import os.path
from tensorflow.python.framework import graph_util
MODEL_DIR = "/home/zheng/PycharmProjects/pbLoad/Models/"
MODEL_NAME = "model"
#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.add(_y, 50, 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(os.path.join(MODEL_DIR,MODEL_NAME), "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)
結果輸出
predictions : [ 101.]
Converted 2 variables to const ops.
9 ops in the final graph.
Tensor("predictions:0", shape=(1,), dtype=float32)
并在指定目錄下生成pb檔案模型，儲存了從輸入層到輸出層這個計算過程的計算圖和相關變量的值，我們得到這個模型後傳入一個輸入，既可以得到一個預估的輸出值
pb模型檔案加載
# -*- coding: utf-8 -*-)
from tensorflow.python.platform import gfile
import tensorflow as tf
from numpy.random import RandomState
sess = tf.Session()
with gfile.FastGFile('./Models/model', 'rb') as f:
graph_def = tf.GraphDef()
graph_def.ParseFromString(f.read())
sess.graph.as_default()
tf.import_graph_def(graph_def, name='') # 導入計算圖
# 需要有一個初始化的過程
sess.run(tf.global_variables_initializer())
# 需要先複原變量
sess.run('W1:0')
sess.run('B1:0')
# 輸入
input_x = sess.graph.get_tensor_by_name('input_holder:0')
#input_y = sess.graph.get_tensor_by_name('y-input:0')
op = sess.graph.get_tensor_by_name('predictions:0')
ret = sess.run(op, feed_dict={input_x:[10]})
print(ret)
輸出結果
[ 101.]
我們可以看到結果一緻。
ckpt格式轉pb格式
通過傳入 CKPT 模型的路徑得到模型的圖和變量資料
通過 import_meta_graph 導入模型中的圖
通過 saver.restore 從模型中恢複圖中各個變量的資料
通過 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 = "/home/zheng/PycharmProjects/ckptToPb/model/"
MODEL_NAME = "frozen_model"
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)) #得到目前圖有幾個操作節點
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("/home/zheng/PycharmProjects/ckptLoad/Models/") #模型目錄
注意改變ckpt模型目錄及pb檔案儲存目錄 。
運作結果為
predictions : [ 101.]
Converted 2 variables to const ops.
9 ops in the final graph.
總結：cpkt檔案格式将模型儲存為4個檔案，pb檔案格式為一個。ckpt模型持久化方式将圖結構與權重參數分開儲存，多了模型更多的細節，适合模型訓練階段;而pb持久化方式完成了從輸入到輸出的前向傳播，完成了端到端的形式，更是個離線使用。
以上這篇tensorflow的ckpt及pb模型持久化方式及轉化詳解就是小編分享給大家的全部内容了，希望能給大家一個參考，也希望大家多多支援我們。
本文标題: tensorflow的ckpt及pb模型持久化方式及轉化詳解
本文位址: http://www.cppcns.com/jiaoben/python/300247.html