TensorFlow模型保存和提取的方法

更新時間：2018年03月08日 14:15:12 投稿：lijiao

這篇文章主要為大家詳細介紹了TensorFlow模型保存和提取的方法，具有一定的參考價值，感興趣的小伙伴們可以參考一下

一、TensorFlow模型保存和提取方法

1. TensorFlow通過tf.train.Saver類實現(xiàn)神經(jīng)網(wǎng)絡模型的保存和提取。tf.train.Saver對象saver的save方法將TensorFlow模型保存到指定路徑中，saver.save(sess,"Model/model.ckpt") ，實際在這個文件目錄下會生成4個人文件：

checkpoint文件保存了一個錄下多有的模型文件列表，model.ckpt.meta保存了TensorFlow計算圖的結構信息，model.ckpt保存每個變量的取值，此處文件名的寫入方式會因不同參數(shù)的設置而不同，但加載restore時的文件路徑名是以checkpoint文件中的“model_checkpoint_path”值決定的。

2. 加載這個已保存的TensorFlow模型的方法是saver.restore(sess,"./Model/model.ckpt") ，加載模型的代碼中也要定義TensorFlow計算圖上的所有運算并聲明一個tf.train.Saver類，不同的是加載模型時不需要進行變量的初始化，而是將變量的取值通過保存的模型加載進來，注意加載路徑的寫法。若不希望重復定義計算圖上的運算，可直接加載已經(jīng)持久化的圖，saver =tf.train.import_meta_graph("Model/model.ckpt.meta") 。

3.tf.train.Saver類也支持在保存和加載時給變量重命名，聲明Saver類對象的時候使用一個字典dict重命名變量即可，{"已保存的變量的名稱name": 重命名變量名}，saver = tf.train.Saver({"v1":u1, "v2": u2})即原來名稱name為v1的變量現(xiàn)在加載到變量u1（名稱name為other-v1）中。

4. 上一條做的目的之一就是方便使用變量的滑動平均值。如果在加載模型時直接將影子變量映射到變量自身，則在使用訓練好的模型時就不需要再調(diào)用函數(shù)來獲取變量的滑動平均值了。載入時，聲明Saver類對象時通過一個字典將滑動平均值直接加載到新的變量中，saver = tf.train.Saver({"v/ExponentialMovingAverage": v})，另通過tf.train.ExponentialMovingAverage的variables_to_restore()函數(shù)獲取變量重命名字典。

此外，通過convert_variables_to_constants函數(shù)將計算圖中的變量及其取值通過常量的方式保存于一個文件中。

二、TensorFlow程序實現(xiàn)

# 本文件程序為配合教材及學習進度漸進進行，請按照注釋分段執(zhí)行 
# 執(zhí)行時要注意IDE的當前工作過路徑，最好每段重啟控制器一次，輸出結果更準確 
 
 
# Part1: 通過tf.train.Saver類實現(xiàn)保存和載入神經(jīng)網(wǎng)絡模型 
 
# 執(zhí)行本段程序時注意當前的工作路徑 
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: 
 sess.run(tf.global_variables_initializer()) 
 saver.save(sess, "Model/model.ckpt") 
 
 
# Part2: 加載TensorFlow模型的方法 
 
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.restore(sess, "./Model/model.ckpt") # 注意此處路徑前添加"./" 
 print(sess.run(result)) # [ 3.] 
 
 
# Part3: 若不希望重復定義計算圖上的運算，可直接加載已經(jīng)持久化的圖 
 
import tensorflow as tf 
 
saver = tf.train.import_meta_graph("Model/model.ckpt.meta") 
 
with tf.Session() as sess: 
 saver.restore(sess, "./Model/model.ckpt") # 注意路徑寫法 
 print(sess.run(tf.get_default_graph().get_tensor_by_name("add:0"))) # [ 3.] 
 
 
# Part4： tf.train.Saver類也支持在保存和加載時給變量重命名 
 
import tensorflow as tf 
 
# 聲明的變量名稱name與已保存的模型中的變量名稱name不一致 
u1 = tf.Variable(tf.constant(1.0, shape=[1]), name="other-v1") 
u2 = tf.Variable(tf.constant(2.0, shape=[1]), name="other-v2") 
result = u1 + u2 
 
# 若直接生命Saver類對象，會報錯變量找不到 
# 使用一個字典dict重命名變量即可，{"已保存的變量的名稱name": 重命名變量名} 
# 原來名稱name為v1的變量現(xiàn)在加載到變量u1（名稱name為other-v1）中 
saver = tf.train.Saver({"v1": u1, "v2": u2}) 
 
with tf.Session() as sess: 
 saver.restore(sess, "./Model/model.ckpt") 
 print(sess.run(result)) # [ 3.] 
 
 
# Part5: 保存滑動平均模型 
 
import tensorflow as tf 
 
v = tf.Variable(0, dtype=tf.float32, name="v") 
for variables in tf.global_variables(): 
 print(variables.name) # v:0 
 
ema = tf.train.ExponentialMovingAverage(0.99) 
maintain_averages_op = ema.apply(tf.global_variables()) 
for variables in tf.global_variables(): 
 print(variables.name) # v:0 
       # v/ExponentialMovingAverage:0 
 
saver = tf.train.Saver() 
 
with tf.Session() as sess: 
 sess.run(tf.global_variables_initializer()) 
 sess.run(tf.assign(v, 10)) 
 sess.run(maintain_averages_op) 
 saver.save(sess, "Model/model_ema.ckpt") 
 print(sess.run([v, ema.average(v)])) # [10.0, 0.099999905] 
 
 
# Part6: 通過變量重命名直接讀取變量的滑動平均值 
 
import tensorflow as tf 
 
v = tf.Variable(0, dtype=tf.float32, name="v") 
saver = tf.train.Saver({"v/ExponentialMovingAverage": v}) 
 
with tf.Session() as sess: 
 saver.restore(sess, "./Model/model_ema.ckpt") 
 print(sess.run(v)) # 0.0999999 
 
 
# Part7: 通過tf.train.ExponentialMovingAverage的variables_to_restore()函數(shù)獲取變量重命名字典 
 
import tensorflow as tf 
 
v = tf.Variable(0, dtype=tf.float32, name="v") 
# 注意此處的變量名稱name一定要與已保存的變量名稱一致 
ema = tf.train.ExponentialMovingAverage(0.99) 
print(ema.variables_to_restore()) 
# {'v/ExponentialMovingAverage': <tf.Variable 'v:0' shape=() dtype=float32_ref>} 
# 此處的v取自上面變量v的名稱name="v" 
 
saver = tf.train.Saver(ema.variables_to_restore()) 
 
with tf.Session() as sess: 
 saver.restore(sess, "./Model/model_ema.ckpt") 
 print(sess.run(v)) # 0.0999999 
 
 
# Part8: 通過convert_variables_to_constants函數(shù)將計算圖中的變量及其取值通過常量的方式保存于一個文件中 
 
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 
 
with tf.Session() as sess: 
 sess.run(tf.global_variables_initializer()) 
 # 導出當前計算圖的GraphDef部分，即從輸入層到輸出層的計算過程部分 
 graph_def = tf.get_default_graph().as_graph_def() 
 output_graph_def = graph_util.convert_variables_to_constants(sess, 
              graph_def, ['add']) 
 
 with tf.gfile.GFile("Model/combined_model.pb", 'wb') as f: 
  f.write(output_graph_def.SerializeToString()) 
 
 
# Part9: 載入包含變量及其取值的模型 
 
import tensorflow as tf 
from tensorflow.python.platform import gfile 
 
with tf.Session() as sess: 
 model_filename = "Model/combined_model.pb" 
 with gfile.FastGFile(model_filename, 'rb') as f: 
  graph_def = tf.GraphDef() 
  graph_def.ParseFromString(f.read()) 
 
 result = tf.import_graph_def(graph_def, return_elements=["add:0"]) 
 print(sess.run(result)) # [array([ 3.], dtype=float32)]

以上就是本文的全部內(nèi)容，希望對大家的學習有所幫助，也希望大家多多支持腳本之家。

您可能感興趣的文章: