Tensorflow:模型變量保存
覺得有用的話,歡迎一起討論相互學習~
參考文獻Tensorflow實戰Google深度學習框架
實驗平台:
Tensorflow1.4.0
python3.5.0
Tensorflow常用保存模型方法
import tensorflow as tf
saver = tf.train.Saver() # 創建保存器
with tf.Session() as sess:
saver.save(sess,"/path/model.ckpt") #保存模型到相應ckpt文件
saver.restore(sess,"/path/model.ckpt") #從相應ckpt文件中恢復模型變量
- 使用tf.train.Saver會保存運行Tensorflow程序所需要的全部信息,然而有時並不需要某些信息。比如在測試或離線預測時,只需要知道如何從神經網絡的輸入層經過前向傳播計算得到輸出層即可,而不需要類似的變量初始化,模型保存等輔助節點的信息。Tensorflow提供了convert_varibales_to_constants函數,通過這個函數可以將計算圖中的變量及其取值通過常量的方式保存,這樣整個Tensorflow計算圖可以統一存放在一個文件中。
將變量取值保存為pb文件
# pb文件保存方法
import tensorflow as tf
from tensorflow.python.framework import graph_util
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.global_variables_initializer()
with tf.Session() as sess:
sess.run(init_op) # 初始化所有變量
# 導出當前計算圖的GraphDef部分,只需要這一部分就可以完成從輸入層到輸出層的計算過程
graph_def = tf.get_default_graph().as_graph_def()
# 將需要保存的add節點名稱傳入參數中,表示將所需的變量轉化為常量保存下來。
output_graph_def = graph_util.convert_variables_to_constants(sess, graph_def, ['add'])
# 將導出的模型存入文件中
with tf.gfile.GFile("Saved_model/combined_model.pb", "wb") as f:
f.write(output_graph_def.SerializeToString())
# 2. 加載pb文件。
from tensorflow.python.platform import gfile
with tf.Session() as sess:
model_filename = "Saved_model/combined_model.pb"
# 讀取保存的模型文件,並將其解析成對應的GraphDef Protocol Buffer
with gfile.FastGFile(model_filename, 'rb') as f:
graph_def = tf.GraphDef()
graph_def.ParseFromString(f.read())
# 將graph_def中保存的圖加載到當前圖中,其中保存的時候保存的是計算節點的名稱,為add
# 但是讀取時使用的是張量的名稱所以是add:0
result = tf.import_graph_def(graph_def, return_elements=["add:0"])
print(sess.run(result))
# Converted 2 variables to const ops.
# [array([3.], dtype=float32)]