tensorflow的ckpt及pb模型持久化方式及转化详解

# 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)) #得到当前图有几个操作节点

# -*- 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]})

# 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)

# -*- 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)

# 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/") #模型目录

predictions : [ 101.]
Converted 2 variables to const ops.
9 ops in the final graph.