Tensorflow模型持久化的示例分析

這篇文章給大家分享的是有關(guān)Tensorflow模型持久化的示例分析的內(nèi)容。小編覺得挺實(shí)用的，因此分享給大家做個(gè)參考，一起跟隨小編過來看看吧。

當(dāng)我們使用 tensorflow 訓(xùn)練神經(jīng)網(wǎng)絡(luò)的時(shí)候，模型持久化對(duì)于我們的訓(xùn)練有很重要的作用。

如果我們的神經(jīng)網(wǎng)絡(luò)比較復(fù)雜，訓(xùn)練數(shù)據(jù)比較多，那么我們的模型訓(xùn)練就會(huì)耗時(shí)很長(zhǎng)，如果在訓(xùn)練過程中出現(xiàn)某些不可預(yù)計(jì)的錯(cuò)誤，導(dǎo)致我們的訓(xùn)練意外終止，那么我們將會(huì)前功盡棄。為了避免這個(gè)問題，我們就可以通過模型持久化（保存為CKPT格式）來暫存我們訓(xùn)練過程中的臨時(shí)數(shù)據(jù)。

如果我們訓(xùn)練的模型需要提供給用戶做離線的預(yù)測(cè)，那么我們只需要前向傳播的過程，只需得到預(yù)測(cè)值就可以了，這個(gè)時(shí)候我們就可以通過模型持久化（保存為PB格式）只保存前向傳播中需要的變量并將變量的值固定下來，這個(gè)時(shí)候只需用戶提供一個(gè)輸入，我們就可以通過模型得到一個(gè)輸出給用戶。

保存為 CKPT 格式的模型

定義運(yùn)算過程

聲明并得到一個(gè) Saver

通過 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): #創(chuàng)建目錄
  tf.gfile.MakeDirs(MODEL_DIR)

#下面的過程你可以替換成CNN、RNN等你想做的訓(xùn)練過程，這里只是簡(jiǎn)單的一個(gè)計(jì)算公式
input_holder = tf.placeholder(tf.float32, shape=[1], name="input_holder") #輸入占位符，并指定名字，后續(xù)模型讀取可能會(huì)用的
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") #輸出節(jié)點(diǎn)名字，后續(xù)模型讀取會(huì)用到，比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]}) #輸入一個(gè)數(shù)據(jù)測(cè)試一下
  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)) #得到當(dāng)前圖有幾個(gè)操作節(jié)點(diǎn)

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

運(yùn)行后生成的文件如下：

checkpoint ： 記錄目錄下所有模型文件列表
ckpt.data ： 保存模型中每個(gè)變量的取值
ckpt.meta ： 保存整個(gè)計(jì)算圖的結(jié)構(gòu)

保存為 PB 格式模型

定義運(yùn)算過程
通過 get_default_graph().as_graph_def() 得到當(dāng)前圖的計(jì)算節(jié)點(diǎn)信息
通過 graph_util.convert_variables_to_constants 將相關(guān)節(jié)點(diǎn)的values固定
通過 tf.gfile.GFile 進(jìn)行模型持久化

# 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): #創(chuàng)建目錄
#   tf.gfile.MakeDirs(MODEL_DIR)

output_graph = "model/pb/add_model.pb"

#下面的過程你可以替換成CNN、RNN等你想做的訓(xùn)練過程，這里只是簡(jiǎn)單的一個(gè)計(jì)算公式
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") #做一個(gè)加法運(yùn)算

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() #得到當(dāng)前的圖的 GraphDef 部分，通過這個(gè)部分就可以完成重輸入層到輸出層的計(jì)算過程

  output_graph_def = graph_util.convert_variables_to_constants( # 模型持久化，將變量值固定
    sess,
    graph_def,
    ["predictions"] #需要保存節(jié)點(diǎn)的名字
  )
  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(): 打印模型節(jié)點(diǎn)信息
#   print (op.name)

*GraphDef：這個(gè)屬性記錄了tensorflow計(jì)算圖上節(jié)點(diǎn)的信息。

add_model.pb ： 里面保存了重輸入層到輸出層這個(gè)計(jì)算過程的計(jì)算圖和相關(guān)變量的值，我們得到這個(gè)模型后傳入一個(gè)輸入，既可以得到一個(gè)預(yù)估的輸出值

CKPT 轉(zhuǎn)換成 PB格式

通過傳入 CKPT 模型的路徑得到模型的圖和變量數(shù)據(jù)
通過 import_meta_graph 導(dǎo)入模型中的圖
通過 saver.restore 從模型中恢復(fù)圖中各個(gè)變量的數(shù)據(jù)
通過 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): #創(chuàng)建目錄
  tf.gfile.MakeDirs(MODEL_DIR)

def freeze_graph(model_folder):
  checkpoint = tf.train.get_checkpoint_state(model_folder) #檢查目錄下ckpt文件狀態(tài)是否可用
  input_checkpoint = checkpoint.model_checkpoint_path #得ckpt文件路徑
  output_graph = os.path.join(MODEL_DIR, MODEL_NAME) #PB模型保存路徑

  output_node_names = "predictions" #原模型輸出操作節(jié)點(diǎn)的名字
  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() #獲得默認(rèn)的圖
  input_graph_def = graph.as_graph_def() #返回一個(gè)序列化的圖代表當(dāng)前的圖

  with tf.Session() as sess:
    saver.restore(sess, input_checkpoint) #恢復(fù)圖并得到數(shù)據(jù)

    print "predictions : ", sess.run("predictions:0", feed_dict={"input_holder:0": [10.0]}) # 測(cè)試讀出來的模型是否正確，注意這里傳入的是輸出 和輸入 節(jié)點(diǎn)的 tensor的名字，不是操作節(jié)點(diǎn)的名字

    output_graph_def = graph_util.convert_variables_to_constants( #模型持久化，將變量值固定
      sess,
      input_graph_def,
      output_node_names.split(",") #如果有多個(gè)輸出節(jié)點(diǎn)，以逗號(hào)隔開
    )
    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)) #得到當(dāng)前圖有幾個(gè)操作節(jié)點(diǎn)

    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是輸入的類型，
  # 這里運(yùn)行程序時(shí)需要帶上模型ckpt的路徑，不然會(huì)報(bào) error: too few arguments
  aggs = parser.parse_args()
  freeze_graph(aggs.model_folder)
  # freeze_graph("model/ckpt") #模型目錄

感謝各位的閱讀！關(guān)于“Tensorflow模型持久化的示例分析”這篇文章就分享到這里了，希望以上內(nèi)容可以對(duì)大家有一定的幫助，讓大家可以學(xué)到更多知識(shí)，如果覺得文章不錯(cuò)，可以把它分享出去讓更多的人看到吧！