tensorflow模型文件(ckpt)转pb文件的方法（不知道输出节点名）

tensorflow模型文件模型文件(ckpt)转转pb文件的方法（不知道输出节点名）文件的方法（不知道输出节点名）

主要介绍了tensorflow模型文件(ckpt)转pb文件（不知道输出节点名），本文通过实例代码给大家介绍的非常详细，对大家的学习或工作具有一定的参考借鉴价值,需要的朋友可以

参考下

网上关于tensorflow模型文件ckpt格式转pb文件的帖子很多，本人几乎尝试了所有方法，最后终于成功了，现总结如下。方法无外乎下面两种：

使用tensorflow.python.tools.freeze_graph.freeze_graph

使用graph_util.convert_variables_to_constants

1、tensorflow模型的文件解读

使用tensorflow训练好的模型会自动保存为四个文件，如下

2、ckpt转pb文件（freeze_graph.freeze_graph）

此种方法尝试成功，虽然不知道输出节点名，但是只要模型代码还在就可以操作，直接上代码。

import tensorflow as tf

import os

from tensorflow.python.tools import freeze_graph

from model import network # network是你们自己定义的模型结构（代码结构）

# egs:

# def network(input)：

# return tf.layers.softmax(input)

output_node = network(input_node) # 神经网络的输出

# 设置输出数据类型（特别注意，这里必须要跟输出网络参数的数据格式保持一致，不然会导致模型预测 精度或者预测能力的丢失）以及重新定义输出节点的名字（这样在后面保存pb文件以及之后使用pb文件时直接使用重新定义的节点名字即可）

flow = tf.cast(output_node , tf.float16, 'the_outputs')

saver = tf.train.Saver()

with tf.Session() as sess:

saver.restore(sess, model_path)

#保存模型图（结构），为一个json文件

tf.train.write_graph(sess.graph_def, 'output_model/pb_model', 'model.pb')

#将模型参数与模型图结合，并保存为pb文件

freeze_graph.freeze_graph('output_model/pb_model/model.pb', '', False, model_path, 'the_outputs','save/restore_all', 'save/Const:0', 'output_model/pb_model/frozen_model.pb', False, "")

'''

:param input_checkpoint:

:param output_graph: PB模型保存路径

:return:

'''

# checkpoint = tf.train.get_checkpoint_state(model_folder) #检查目录下ckpt文件状态是否可用

# input_checkpoint = checkpoint.model_checkpoint_path #得ckpt文件路径

# 指定输出的节点名称,该节点名称必须是原模型中存在的节点

output_node_names = "InceptionV3/Logits/SpatialSqueeze"

output_graph_def = graph_util.convert_variables_to_constants( # 模型持久化，将变量值固定

sess=sess,

input_graph_def=input_graph_def,# 等于:sess.graph_def

output_node_names=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():

out_pb_path="models/pb/frozen_model.pb"

# 调用freeze_graph将ckpt转为pb

freeze_graph(input_checkpoint,out_pb_path)

参考链接：

https://www.jb51.net/article/185209.htm

https://www.jb51.net/article/185206.htm

到此这篇关于tensorflow模型文件(ckpt)转pb文件（不知道输出节点名）的文章就介绍到这了,更多相关tensorflow ckpt转pb文件内容请搜索我们以前的文章或继续浏览下面的相关文章希望大家以后多多

支持我们！