tensorflow模型固化

1 tensorflow模型固化 1.1 训练时直接固化成pb文件

import tensorflow as tf
from tensorflow.python.framework import graph_util

v1 = tf.Variable(tf.constant(1.0, tf.float32, [1]),name='v1')
v2 = tf.Variable(tf.constant(2.0, tf.float32, [1]),name='v2')
result = v1 + v2

with tf.Session() as sess:
    print(sess.run(tf.global_variables_initializer()))
    print(result.eval())
    # 导出当前计算图的GraphDef部分,只需要这一部分就可以完成从输入层到输出层的计算过程
    graph_def = tf.get_default_graph().as_graph_def()

    # 将图中的变量及其取值转化为常量,同时将图中不必要的节点去掉(例如一些诸如变量初始化 *** 作的系统运算)
    # 如果只关心程序中定义的某些运算时,和这些计算无关的节点就没有必要导出并保存了,在下面的一行代码中,
    # 最后一个参数['add']给出了需要保存的节点名称.add节点是上面定义的两个变量相加的 *** 作.
    # 注意这里给出的计算节点的名称,所以没有后面的:0,:0表示的是该节点的第一个输出
    output_graph_def = graph_util.convert_variables_to_constants(sess,graph_def,['add'])

    # 将导出的模型存入文件
    with tf.gfile.GFile('./model/combined_model.pb', 'wb') as f:
        f.write(output_graph_def.SerializeToString())

1.2 已有训练好的ckpt文件固化成pb文件 1.2.1 前期准备

import tensorflow as tf

v1 = tf.Variable(tf.constant(1.0,tf.float32, [1]),name='v1')
v2 = tf.Variable(tf.constant(2.0,tf.float32, [1]),name='v2')

result = v1 + v2
saver = tf.train.Saver()

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    print(result.eval())
    saver.save(sess,'./model/model.ckpt')

模型文件如下

整个模型固化过程分两步进行，首先加载模型并将参数赋予训练得到的值，然后固化模型，写入pb文件。

1.2.2 加载模型和参数

方式一，通过加载meta图结构加载原图， saver = tf.train.import_meta_graph(cpkt_path + ‘.meta’, clear_devices=True)

import tensorflow as tf
from tensorflow.python.framework import graph_util

saver = tf.train.import_meta_graph('./FAW_VHE_best_model/ckpt.meta')
graph_def = tf.get_default_graph().as_graph_def()

with tf.Session() as sess:
    saver.restore(sess, './FAW_VHE_best_model/ckpt')
    output_graph_def = graph_util.convert_variables_to_constants(sess,graph_def,['add'])
    print(sess.run(tf.get_default_graph().get_tensor_by_name('add:0')))  # 3.0
    with tf.gfile.GFile('./model/model.pb', 'wb') as f:
        f.write(output_graph_def.SerializeToString())  # 得到文件：model.pb

————————————————
原文链接：https://blog.csdn.net/liugan528/article/details/109482836

方法二 通过重新构建与训练模型一样的图，初始化参数，再读取模型文件将参数赋值给模型参数上

def export_model_test():
	#测试placeholder能不能在保存pb文件的时候重命名
	model_path = "./FAW_VHE_best_model/ckpt"
	tf.reset_default_graph()   #
	model = Model(user_num, poi_num, p_adj)
	model.build()
	# model.u = tf.placeholder(tf.float32, shape=[None], name='u')
	# model.hist = tf.placeholder(tf.float32, shape=[None, None], name='hist')
	# model.length = tf.placeholder(tf.float32, shape=[None], name='length')
	# # max_length = tf.placeholder(tf.float32, shape=[None, 224, 224, 3], name='input')
	# model.pos = tf.placeholder(tf.float32, shape=[None], name='pos')
	# model.neg = tf.placeholder(tf.float32, shape=[None], name='neg')
	# tf.identity(model.neg, "neg")
	tf.identity(model.u, "u")   # op可以重命名，并且使用，placeholder不能用这种方式重命名，后面加载pb文件看节点的时候看不到placeholder的重命名结果，op可以看到
	tf.identity(model.hist, "hist")
	tf.identity(model.length, "length")

	tf.identity(model.score, "out")
	saver = tf.train.Saver()

	with tf.Session() as sess:
		sess.run(tf.global_variables_initializer())
		saver.restore(sess, model_path)
		graph_def = tf.get_default_graph().as_graph_def()
		output_graph_def = tf.graph_util.convert_variables_to_constants(
			sess, graph_def, ["out"])
		with tf.gfile.GFile('./model/model.pb', 'wb') as f:
			f.write(output_graph_def.SerializeToString())

1.2.3 固化模型写入pb文件

方式一 把变量转成常量之后写入PB文件中

output_graph_def = graph_util.convert_variables_to_constants(sess,graph_def,['add'])
# 将导出的模型存入文件
with tf.gfile.GFile('./model/model.pb', 'wb') as f:
    f.write(output_graph_def.SerializeToString())

方式二 通过freeze_graph把tf.train.write_graph()生成的pb文件与tf.train.saver()生成的chkp文件固化之后重新生成一个pb文件，这一种现在不太建议使用

import tensorflow as tf
from tensorflow.python.tools import freeze_graph 
 
# network是你自己定义的模型
import network
 
# 模型的checkpoint文件地址
ckpt_path = "./FAW_VHE_best_model/ckpt"
 
def main(): 
	tf.reset_default_graph() 
 
	x = tf.placeholder(tf.float32, shape=[None, 224, 224, 3], name='input')
 
        # flow是模型的输出
	flow = network(x)
        #设置输出类型以及输出的接口名字，为了之后的调用pb的时候使用 
        flow = tf.cast(flow, tf.int8, 'out')
 
	with tf.Session() as sess:
		#保存图，在./pb_model文件夹中生成model.pb文件
                # model.pb文件将作为input_graph给到接下来的freeze_graph函数
		tf.train.write_graph(sess.graph_def, './pb_model', 'model.pb') 
		#把图和参数结构一起 
		freeze_graph.freeze_graph(
			input_graph='./model/model.pb',
			input_saver='',
			input_binary=False, 
			input_checkpoint=ckpt_path, 
			output_node_names='out',
			restore_op_name='save/restore_all',
			filename_tensor_name='save/Const:0',
			output_graph='./pb_model/frozen_model.pb',
			clear_devices=False,
			initializer_nodes=''
			)
 
	print("done") 
 
if __name__ == '__main__': 
	main()

方式 2 读取pb文件 2.1 重新定义placeholder，并通过input_map链接的方式

import tensorflow as tf
from tensorflow.python.platform import gfile

v1_ph = tf.placeholder(tf.float32,name='v1_ph')
v2_ph = tf.placeholder(tf.float32,name='v2_ph')

with tf.Session() as sess:
    model_filename = './model/model.pb'
    # 读取保存的模型文件,并将文件解析成对应的GraphDef Protobuf Buffer
    with gfile.FastGFile(model_filename,'rb') as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())

        # 查看计算图上有哪些节点，用于确定网络的输入输出节点名
        node_names_list = [tensor.name for tensor in graph_def.node]
        for node_name in node_names_list:
            print(node_name)

    # 将graph_def中保存的图加载到当前的图中,return_elements=['add:0']给出了返回的张量名称
    # 在保存的时候给出的是计算节点的名称,所以为'add',在加载的时候给出的是张量的名称,所以是'add:0'
    result = tf.import_graph_def(graph_def, input_map={'v1:0':v1_ph, 'v2:0':v2_ph}, return_elements=['add:0'])
    print(sess.run(result, feed_dict={v1_ph:177, v2_ph:12}))  # 输出189.0

————————————————
原文链接：https://blog.csdn.net/liugan528/article/details/109482836

2.2 placeholder重新定义，直接通过return_elements返回的方式

import tensorflow as tf
from tensorflow.python.platform import gfile
with tf.Session() as sess:
    model_filename = './model/combined_model.pb'
    # 读取保存的模型文件,并将文件解析成对应的GraphDef Protobuf Buffer
    with gfile.FastGFile(model_filename,'rb') as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
    # 将graph_def中保存的图加载到当前的图中,return_elements=['add:0']给出了返回的张量名称
    # 在保存的时候给出的是计算节点的名称,所以为'add',在加载的时候给出的是张量的名称,所以是'add:0'
    v1_ph,v2_ph,result = tf.import_graph_def(graph_def, return_elements=['v1:0', 'v2:0','add:0'])
    print(sess.run(result, feed_dict={v1_ph:177, v2_ph:12}))  # 输出189.0

其他tensorflow有用函数

import tensorflow as tf
from tensorflow.python.framework import graph_util
from tensorflow.python import pywrap_tensorflow
 
def freeze_graph(cpkt_path, pb_path):
    # checkpoint = tf.train.get_checkpoint_state('zhou/') #检查目录下ckpt文件状态是否可用
    # cpkt_path2 = checkpoint.model_checkpoint_path #得ckpt文件路径
    # cpkt_path3 = checkpoint.all_model_checkpoint_paths
    # print("model_pa:",cpkt_path3)
 
    # 指定输出的节点名称,该节点名称必须是原模型中存在的节点
    # output_node_names = "num_detections,raw_detection_boxes,raw_detection_scores"
    output_node_names = 'logistic_loss/mul'
    saver = tf.train.import_meta_graph(cpkt_path + '.meta', clear_devices=True)
    graph = tf.get_default_graph()
    input_graph_def = graph.as_graph_def()
    print(graph)
    # feature_data_list = input_graph_def.get_operation_by_name('resnet_v2_50/conv1').outputs[0]
    # input_image=tf.placeholder(None,28,28,1)
    with tf.Session() as sess:
        saver.restore(sess, cpkt_path)  # 恢复图并得到数据
        pb_path_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(","))  # 如果有多个输出节点，以逗号隔开
        # print(pb_path_def)
 
        with tf.gfile.GFile(pb_path, 'wb') as fgraph:
            fgraph.write(pb_path_def.SerializeToString())