tensorflow对深度学习生成的pb模型文件的保存与读取及节点和张量的输出

一、pb模型的保存

1、MTCNN人脸检测算法中官网训练好的参数保存在三个文件名称分别为&＃xff1a;det1.npy、det2.npy、det3.npy的后缀名为 .npy文件中&＃xff08;.npy文件也是一种以二进制保存的文件&＃xff09;&＃xff0c;将.npy文件转换为 .pb 模型文件的方法通过以下代码实现&＃xff1a;

import tensorflow as tf
import detect_face
import os
from tensorflow.python.framework.graph_util import convert_variables_to_constants# tf.InteractiveSession()让自己成为默认的session&＃xff0c;用户不需指明哪个session运行情况下就可运行起来
# tf.InteractiveSession()来构建会话时&＃xff0c;可以先构建一个session然后再定义操作
# 使用tf.Session()来构建会话&＃xff0c;需要在会话构建之前定义好全部的操作&＃xff08;operation&＃xff09;然后再构建会话
sess &＃61; tf.InteractiveSession()with tf.variable_scope("pnet"):data &＃61; tf.placeholder(tf.float32, (None, None, None, 3), "input")pnet &＃61; detect_face.PNet({"data": data})pnet.load("det1.npy", sess)with tf.variable_scope("rnet"):data &＃61; tf.placeholder(tf.float32, (None, 24, 24, 3), "input")rnet &＃61; detect_face.RNet({"data": data})rnet.load("det2.npy", sess)with tf.variable_scope("onet"):data &＃61; tf.placeholder(tf.float32, (None, 48, 48, 3), "input")onet &＃61; detect_face.ONet({"data": data})onet.load("det3.npy", sess)# 将存储到 .npy文件中的网络模型参数转换成用 .bp文件存储的模型格式
"""
tf模型导出为单个文件&＃xff08;同时包含模型架构定义与权重&＃xff09;
利用tf.train_write_graph()默认情况下只导出了网络的定义&＃xff08;无权重&＃xff09;
利用tf.train_Saver().Save()导出文件graph_def与权重分离的&＃xff0c;graph_def没有包含网络中的Variable值
&＃xff08;通常情况只存储了权重&＃xff09;&＃xff0c;但却包含了constant值&＃xff0c;如果把Variable转换成constant&＃xff0c;
可达到使用一个文件同时存储网络架构与权重的目标
convert_variables_to_constants函数会将计算图中的变量取值以常量的形式保存&＃xff0c;
在保存模型文件的时候只是导出了GraphDef部分&＃xff0c;GraphDef保存了从输入到输出的计算过程
保存的时候通过convert_variables_to_constants函数来指定保存的节点名称而不是张量的名称
比如&＃xff1a;“add:0”是张量的名称&＃xff0c;而"add"表示的是节点的名称。
"""
#
constant_graph &＃61; convert_variables_to_constants(sess, sess.graph_def,["pnet/input", "rnet/input", "onet/input","pnet/conv4-2/BiasAdd", "pnet/prob1","rnet/conv5-2/conv5-2", "rnet/prob1","onet/conv6-2/conv6-2", "onet/conv6-3/conv6-3","onet/prob1"])with tf.gfile.FastGFile("face_detect.pb", mode&＃61;"wb") as f:f.write(constant_graph.SerializeToString())

        以上方法是传统保存pb模型的方法&＃xff0c;如果将pb模型用于TensorFlow serving部署线上模型事&＃xff0c;需要生成一个带有signature的pb模型文件&＃xff0c;应用saved_model模块&＃xff0c;具体代码如下&＃xff08;承接上面代码&＃xff09;&＃xff1a;

"""
saved_model模块主要用于TensorFlow serving&＃xff0c;TF serving是一个将训练好的模型部署
至生产环境的系统&＃xff0c;主要优点在于&＃xff1a;保持server端与API不变情况下部署新的算法或进行
实验&＃xff0c;同时还有很高的性能。
"""
# 构造SavedModelBuilder对象&＃xff0c;初始化方法只需要传入传入用于保存模型的目录名&＃xff0c;目录不用预先创建
builder &＃61; tf.saved_model.builder.SavedModelBuilder("./saved_model_dir_signature/")# 输入&＃xff08;三个&＃xff09;与输出&＃xff08;七个&＃xff09;tensor获取
pnet_input_tensor &＃61; tf.get_default_graph().get_tensor_by_name("pnet/input:0")
rnet_input_tensor &＃61; tf.get_default_graph().get_tensor_by_name("rnet/input:0")
onet_input_tensor &＃61; tf.get_default_graph().get_tensor_by_name("onet/input:0")pnet_conv4_2_BiasAdd_output_tensor &＃61; tf.get_default_graph().get_tensor_by_name("pnet/conv4-2/BiasAdd:0")
pnet_prob1_output_tensor &＃61; tf.get_default_graph().get_tensor_by_name("pnet/prob1:0")
rnet_conv5_2_conv5_2_output_tensor &＃61; tf.get_default_graph().get_tensor_by_name("rnet/conv5-2/conv5-2:0")
rnet_prob1_output_tensor &＃61; tf.get_default_graph().get_tensor_by_name("rnet/prob1:0")
onet_conv6_2_conv6_2_output_tensor &＃61; tf.get_default_graph().get_tensor_by_name("onet/conv6-2/conv6-2:0")
onet_conv6_3_conv6_3_output_tensor &＃61; tf.get_default_graph().get_tensor_by_name("onet/conv6-3/conv6-3:0")
onet_prob1_output_tensor &＃61; tf.get_default_graph().get_tensor_by_name("onet/prob1:0")
"""
tf.saved_model.utils.build_tensor_info的作用是构建一个TensorInfo protocol
输入参数是张量的名称、类型大小&＃xff0c;这里是string&＃xff0c;应是名称&＃xff1b;输出是基于提供参数的a tensor protocol buffer
"""
pnet_inputs &＃61; tf.saved_model.utils.build_tensor_info(pnet_input_tensor)
rnet_inputs &＃61; tf.saved_model.utils.build_tensor_info(rnet_input_tensor)
onet_inputs &＃61; tf.saved_model.utils.build_tensor_info(onet_input_tensor)
pnet_conv4_2_BiasAdd_outputs &＃61; tf.saved_model.utils.build_tensor_info(pnet_conv4_2_BiasAdd_output_tensor)
pnet_prob1_outputs &＃61; tf.saved_model.utils.build_tensor_info(pnet_prob1_output_tensor)
rnet_conv5_2_conv5_2_outputs &＃61; tf.saved_model.utils.build_tensor_info(rnet_conv5_2_conv5_2_output_tensor)
rnet_prob1_outputs &＃61; tf.saved_model.utils.build_tensor_info(rnet_prob1_output_tensor)
onet_conv6_2_conv6_2_outputs &＃61; tf.saved_model.utils.build_tensor_info(onet_conv6_2_conv6_2_output_tensor)
onet_conv6_3_conv6_3_outputs &＃61; tf.saved_model.utils.build_tensor_info(onet_conv6_3_conv6_3_output_tensor)
onet_prob1_outputs &＃61; tf.saved_model.utils.build_tensor_info(onet_prob1_output_tensor)# 整合所有的输入与输出&＃xff0c;都是字典&＃xff0c;key可以自定义&＃xff0c;value是之前build_tensor_info的一个张量协议缓冲区
inputs &＃61; {"pnet_inputs": pnet_inputs,"rnet_inputs": rnet_inputs,"onet_inputs": onet_inputs}
outputs &＃61; {"pnet_conv4_2_BiasAdd_outputs": pnet_conv4_2_BiasAdd_outputs,"pnet_prob1_outputs": pnet_prob1_outputs,"rnet_conv5_2_conv5_2_outputs": rnet_conv5_2_conv5_2_outputs,"rnet_prob1_outputs": rnet_prob1_outputs,"onet_conv6_2_conv6_2_outputs": onet_conv6_2_conv6_2_outputs,"onet_conv6_3_conv6_3_outputs": onet_conv6_3_conv6_3_outputs,"onet_prob1_outputs": onet_prob1_outputs}# 构建signature_def_map
signature &＃61; tf.saved_model.signature_def_utils.build_signature_def(inputs&＃61;inputs, outputs&＃61;outputs, method_name&＃61;"mtcnn_signature"
)
"""
add_meta_graph_and_variables方法导入graph的信息及变量&＃xff0c;该方法假设变量都已经初始化好了&＃xff0c;
对应每个SavedModelBuilder这个方法一定要执行一次用于导入第一个meta_graph,
第一个参数&＃xff1a;传入当前的session&＃xff0c;包含了graph的结构与所有变量&＃xff1b;第二个参数&＃xff1a;给当前需要保存的
meta_graph一个标签&＃xff0c;标签名可以自定义&＃xff0c;在之后模型载入的时候&＃xff0c;需要根据这个签名去查找对应的
MetaGraphDef&＃xff0c;标签也可以选用系统定义好的参数&＃xff0c;如&＃xff1a;tf.saved_model.tag_constants.SERVING
与tf.saved_model.tag_constants.TRAINING
"""
builder.add_meta_graph_and_variables(sess&＃61;sess, tags&＃61;[tf.saved_model.tag_constants.SERVING],signature_def_map&＃61;{"mtcnn_signature": signature})"""
save方法就是模型序列化到指定目录底下,保存好以后到saved_model_dir目录下&＃xff0c;
会有一个saved_model.pb文件以及variable文件夹&＃xff0c;variable文件夹保存所有变量&＃xff0c;
saved_model.pb文件用于保存模型结构等信息
"""
builder.save()sess.close()

二、pb模型的读取及节点和张量的输出

1、普通的pb模型的输出&＃xff0c;代码如下&＃xff1a;

import tensorflow as tf
with tf.Graph().as_default():output_graph_def &＃61; tf.GraphDef()output_graph_path &＃61; &＃39;./20180402-114759.pb&＃39;#sess.graph.add_to_collection("input", mnist.test.images)with open(output_graph_path, "rb") as f:output_graph_def.ParseFromString(f.read())print(output_graph_def) # 输出图的节点操作_ &＃61; tf.import_graph_def(output_graph_def, name&＃61;"")with tf.Session() as sess:tf.global_variables_initializer()test &＃61; sess.graph.get_tensor_by_name("embeddings:0")print(test) # 输出张量操作# Tensor("embeddings:0", shape&＃61;(?, 512), dtype&＃61;float32)

2、带有signature的pb模型文件的读取&＃xff0c;代码如下&＃xff1a;

sess &＃61; tf.InteractiveSession()
saved_model_dir &＃61; "saved_model_dir_signature"
meta_graph_def &＃61; tf.saved_model.loader.load(sess, [tf.saved_model.tag_constants.SERVING], saved_model_dir)
print(meta_graph_def)

3、带有signature的pb模型文件在MTCNN人脸检测算法的使用&＃xff08;其中的一个函数&＃xff09;&＃xff0c;代码如下&＃xff1a;

def create_mtcnn(sess, saved_model_dir):# 载入和使用带有signature的模型meta_graph_def &＃61; tf.saved_model.loader.load(sess, [tf.saved_model.tag_constants.SERVING], saved_model_dir)# 从meta_graph_def中提取signatureDef对象signature &＃61; meta_graph_def.signature_def# 从signature中找出具体输入和输出的对象pnet_input_tensor_name &＃61; signature[signature_key].inputs[pnet_input_key].namernet_input_tensor_name &＃61; signature[signature_key].inputs[rnet_input_key].nameonet_input_tensor_name &＃61; signature[signature_key].inputs[onet_input_key].namepnet_conv4_2_BiasAdd_output_tensor_name &＃61; signature[signature_key].outputs[pnet_conv4_2_BiasAdd_output_key].namepnet_prob1_output_tensor_name &＃61; signature[signature_key].outputs[pnet_prob1_output_key].namernet_conv5_2_conv5_2_output_tensor_name &＃61; signature[signature_key].outputs[rnet_conv5_2_conv5_2_output_key].namernet_prob1_output_tensor_name &＃61; signature[signature_key].outputs[rnet_prob1_output_key].nameonet_conv6_2_conv6_2_output_tensor_name &＃61; signature[signature_key].outputs[onet_conv6_2_conv6_2_output_key].nameonet_conv6_3_conv6_3_output_tensor_name &＃61; signature[signature_key].outputs[onet_conv6_3_conv6_3_output_key].nameonet_prob1_output_tensor_name &＃61; signature[signature_key].outputs[onet_prob1_output_key].name# 获取tensor并inferencepnet_input &＃61; sess.graph.get_tensor_by_name(pnet_input_tensor_name)rnet_input &＃61; sess.graph.get_tensor_by_name(rnet_input_tensor_name)onet_input &＃61; sess.graph.get_tensor_by_name(onet_input_tensor_name)pnet_conv4_2_BiasAdd_output &＃61; sess.graph.get_tensor_by_name(pnet_conv4_2_BiasAdd_output_tensor_name)pnet_prob1_output &＃61; sess.graph.get_tensor_by_name(pnet_prob1_output_tensor_name)rnet_conv5_2_conv5_2_output &＃61; sess.graph.get_tensor_by_name(rnet_conv5_2_conv5_2_output_tensor_name)rnet_prob1_output &＃61; sess.graph.get_tensor_by_name(rnet_prob1_output_tensor_name)onet_conv6_2_conv6_2_output &＃61; sess.graph.get_tensor_by_name(onet_conv6_2_conv6_2_output_tensor_name)onet_conv6_3_conv6_3_output &＃61; sess.graph.get_tensor_by_name(onet_conv6_3_conv6_3_output_tensor_name)onet_prob1_output &＃61; sess.graph.get_tensor_by_name(onet_prob1_output_tensor_name)# lambda:匿名函数pnet_fun &＃61; lambda img: sess.run((pnet_conv4_2_BiasAdd_output, pnet_prob1_output), feed_dict&＃61;{pnet_input: img})rnet_fun &＃61; lambda img: sess.run((rnet_conv5_2_conv5_2_output, rnet_prob1_output), feed_dict&＃61;{rnet_input: img})onet_fun &＃61; lambda img: sess.run((onet_conv6_2_conv6_2_output, onet_conv6_3_conv6_3_output, onet_prob1_output),feed_dict&＃61;{onet_input: img})return pnet_fun, rnet_fun, onet_fun