生成ckpt格式的模型
# 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): #创建目录
tf.gfile.MakeDirs(MODEL_DIR)
#下面的过程你可以替换成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.greater(_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)) #得到当前图有几个操作节点
for op in tf.get_default_graph().get_operations(): #打印模型节点信息
print(op.name, op.values())
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 = "model/pb"
MODEL_NAME = "frozen_model.pb"
if not tf.gfile.Exists(MODEL_DIR): #创建目录
tf.gfile.MakeDirs(MODEL_DIR)
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)) #得到当前图有几个操作节点
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是输入的类型,
# 这里运行程序时需要带上模型ckpt的路径,不然会报 error: too few arguments
aggs = parser.parse_args()
freeze_graph(aggs.model_folder)
# freeze_graph("model/ckpt") #模型目录
在pycharm下调试Alt+Shift+F9,然后按0,进入参数设置界面,参数设置如下(请根据自己具体路径设置):
在pycharm 命令窗口执行:
(py35) D:\documents\pycharm\daily>python convert_ckpt_to_pb.py .\model\ckpt
生成的pb文件用Netron生成一下网络结构图
生成pb模型是为了用来预测,下面说明pb文件的使用。
import tensorflow as tf
import os
from tensorflow.python.platform import gfile
pb_file_path = os.getcwd() #获取当前代码路径
pb_file_path += "\\model\\pb\\"
sess = tf.Session()
with gfile.FastGFile(pb_file_path + 'frozen_model.pb', '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())
#下面三句,是能否复现模型的关键
# 输入
input_x = sess.graph.get_tensor_by_name("input_holder:0") #此处的x一定要和之前保存时输入的名称一致!
#input_y = sess.graph.get_tensor_by_name('y:0') #此处的y一定要和之前保存时输入的名称一致!
op = sess.graph.get_tensor_by_name("predictions:0") #此处的op_to_store一定要和之前保存时输出的名称一致!
ret = sess.run(op, feed_dict={input_x: [5.0]})
print(ret)
# 输出 False
这里测试注意名字要一致,输入为5.0,进入计算图,5×5+1=26<50返回False,即输出False。