tensorflow的ckpt及pb模型持久化方式及轉化詳解
使用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,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,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 : ",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,name="input_holder") W1 = tf.Variable(tf.constant(5.0,name="predictions") init = tf.global_variables_initializer() with tf.Session() as sess: sess.run(init) print "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模型持久化方式及轉化詳解就是小編分享給大家的全部內容了,希望能給大家一個參考,也希望大家多多支援我們。