tensorflow的ckpt及pb模型持久化方式及转化详解
使用tensorflow训练模型的时候，模型持久化对我们来说非常重要。
如果我们的模型比较复杂，需要的数据比较多，那么在模型的训练时间会耗时很长。如果在训练过程中出现了模型不可预期的错误，导致训练意外终止，那么我们将会前功尽弃。为了解决这一问题，我们可以使用模型持久化（保存为ckpt文件格式）来保存我们在训练过程中的临时数据。、
如果我们训练出的模型需要提供给用户做离线预测，那么我们只需要完成前向传播过程。这个时候我们就可以使用模型持久化（保存为pb文件格式）来只保存前向传播过程中的变量并将变量固定下来，这时候用户只需要提供一个输入即可得到前向传播的预测结果。
ckpt和pb持久化方式的区别在于ckpt文件将模型结构与模型权重分离保存，便于训练过程;pb文件则是graph_def的序列化文件，便于发布和离线预测。官方提供freeze_grpah.py脚本来将ckpt文件转为pb文件。
CKPT模型持久化
首先定义前向传播过程;
声明并得到一个Saver;
使用Saver.save()保存模型;
#coding=UTF-8支持中文编码格式
importtensorflowastf
importshutil
importos.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用于保存模型
withtf.Session()assess:
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("%dopsinthefinalgraph."%len(tf.get_default_graph().as_graph_def().node))#得到当前图有几个操作节点
predictions:[101.]
28opsinthefinalgraph.
注：代码含义请参考注释，需要注意的是可以自定义模型保存的路径
ckpt模型持久化使用起来非常简单，只需要我们声明一个tf.train.Saver，然后调用save()函数，将会话模型保存到指定的目录。执行代码结果，会在我们指定模型目录下出现4个文件
checkpoint：记录目录下所有模型文件列表
ckpt.data：保存模型中每个变量的取值
ckpt.meta：保存整个计算图的结构
ckpt模型加载
#-*-coding:utf-8-*-)
importtensorflowastf
fromnumpy.randomimportRandomState
#定义训练数据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()
#creareasession，创建一个会话来运行TensorFlow程序
withtf.Session()assess:
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]})
代码结果，可以看到运行结果一样
predictions:[101.]
PB模型持久化
定义运算过程
通过get_default_graph().as_graph_def()得到当前图的计算节点信息
通过graph_util.convert_variables_to_constants将相关节点的values固定
通过tf.gfile.GFile进行模型持久化
#coding=UTF-8
importtensorflowastf
importshutil
importos.path
fromtensorflow.python.frameworkimportgraph_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()
withtf.Session()assess:
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"]#需要保存节点的名字
)
withtf.gfile.GFile(os.path.join(MODEL_DIR,MODEL_NAME),"wb")asf:#保存模型
f.write(output_graph_def.SerializeToString())#序列化输出
print("%dopsinthefinalgraph."%len(output_graph_def.node))
print(predictions)
#foropintf.get_default_graph().get_operations():打印模型节点信息
#print(op.name)
结果输出
predictions:[101.]
Converted2variablestoconstops.
9opsinthefinalgraph.
Tensor("predictions:0",shape=(1,),dtype=float32)
并在指定目录下生成pb文件模型，保存了从输入层到输出层这个计算过程的计算图和相关变量的值，我们得到这个模型后传入一个输入，既可以得到一个预估的输出值
pb模型文件加载
#-*-coding:utf-8-*-)
fromtensorflow.python.platformimportgfile
importtensorflowastf
fromnumpy.randomimportRandomState
sess=tf.Session()
withgfile.FastGFile('./Models/model','rb')asf:
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)
输出结果
[101.]
我们可以看到结果一致。
ckpt格式转pb格式
通过传入CKPT模型的路径得到模型的图和变量数据
通过import_meta_graph导入模型中的图
通过saver.restore从模型中恢复图中各个变量的数据
通过graph_util.convert_variables_to_constants将模型持久化
#coding=UTF-8
importtensorflowastf
importos.path
importargparse
fromtensorflow.python.frameworkimportgraph_util
MODEL_DIR="/home/zheng/PycharmProjects/ckptToPb/model/"
MODEL_NAME="frozen_model"
deffreeze_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：Whetherornottoclearthedevicefieldforan`Operation`or`Tensor`duringimport.
graph=tf.get_default_graph()#获得默认的图
input_graph_def=graph.as_graph_def()#返回一个序列化的图代表当前的图
withtf.Session()assess:
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(",")#如果有多个输出节点，以逗号隔开
)
withtf.gfile.GFile(output_graph,"wb")asf:#保存模型
f.write(output_graph_def.SerializeToString())#序列化输出
print("%dopsinthefinalgraph."%len(output_graph_def.node))#得到当前图有几个操作节点
if__name__=='__main__':
#parser=argparse.ArgumentParser()
#parser.add_argument("model_folder",type=str,help="inputckptmodeldir")#命令行解析，help是提示符，type是输入的类型，
#这里运行程序时需要带上模型ckpt的路径，不然会报error:toofewarguments
#aggs=parser.parse_args()
#freeze_graph(aggs.model_folder)
freeze_graph("/home/zheng/PycharmProjects/ckptLoad/Models/")#模型目录
注意改变ckpt模型目录及pb文件保存目录。
运行结果为
predictions:[101.] Converted2variablestoconstops. 9opsinthefinalgraph.
总结：cpkt文件格式将模型保存为4个文件，pb文件格式为一个。ckpt模型持久化方式将图结构与权重参数分开保存，多了模型更多的细节，适合模型训练阶段;而pb持久化方式完成了从输入到输出的前向传播，完成了端到端的形式，更是个离线使用。
以上这篇tensorflow的ckpt及pb模型持久化方式及转化详解就是小编分享给大家的全部内容了，希望能给大家一个参考，也希望大家多多支持毛票票。
声明：本文内容来源于网络，版权归原作者所有，内容由互联网用户自发贡献自行上传，本网站不拥有所有权，未作人工编辑处理，也不承担相关法律责任。如果您发现有涉嫌版权的内容，欢迎发送邮件至：czq8825#qq.com（发邮件时，请将#更换为@）进行举报，并提供相关证据，一经查实，本站将立刻删除涉嫌侵权内容。