TensorFlow框架里的各种模型们
1：前言
提起TensorFlow的模型，大家最熟知的莫过于checkpoint文件了，但是其实TensorFlow 1.0 以及2.0 提供了多种不同的模型导出格式，除了checkpoint文件，TensorFlow2.0官方推荐SavedModel格式，使用tf.serving部署模型的时候采用的就是它，此外还有Keras model（HDF5）、Frozen GraphDef，以及用于移动端，嵌入式的TFLite。
本文主要介绍tf.serving以及Tensor-RT依赖的俩种模型结构——SavedModel以及Frozen GraphDef，帮助大家搞清楚TensorFlow到底拥哪些类型的模型，模型与模型之间又有怎样的区别，以及其适应的各种使用场景。
2：模型概览
在开始介绍之前，我们需要明确的是，TensorFlow最核心的思想是：所有数据的计算都利用计算图（computational graph，简称Graph）的方式来表达，TensorFlow官网有这样的一句介绍：“A computational graph is a series of TensorFlow operations arranged into a graph of nodes”，而每个graph中的node也被称为op（即operation），它可以是卷积操作，也可以是简单的数学操作，也可以是Variables和Constants。
首先，模型的导出主要包含了：参数以及网络结构的导出，不同的导出格式可能是分别导出，或者是整合成一个独立的文件，大概可以分为以下三种：
-
参数和网络结构分开保存：checkpoint， SavedModel
-
只保存权重：HDF5（可选）
-
参数和网络结构保存在一个文件：Frozen GraphDef，HDF5（可选）
如果模型很复杂，我们需要对模型结构有一个较为直观的认识，那么TensorBoard——TensorFlow模型可视化工具就比不可少了：
Fig. 1. TensorFlow图的可视化 (Source: TensorFlow website)
除了计算图（computational graph）的概念，TensorFlow里还有很多别的"图"，比如MetaGraph，GraphDef和Frozen GraphDef，不要担心，几行代码带你搞清楚他们。
2.1： MetaGraph简介
首先从大家最熟悉的tf.train.Saver()/saver.restore()开始，其保存好的模型文件结构如下所示：
saved_model/
├── checkpoint
├── model.data-00000-of-00001
├── model.index
└── model.meta
其中Checkpoint 记录了模型文件的保存信息；MetaGraph记录计算图中节点的信息以及运行计算图中节点所需要的元数据，MetaGraph是由Protocol Buffer定义的MetaGraphDef保存在.meta文件中；其中模型经过训练的模型参数，权重，可训练的变量保存在.data文件中；张量名到张量的对应映射关系保存在.index文件中。
从 Meta Graph 中恢复构建的图包含 Variable 的信息，但却没有 Variable 的实际值，所以， 从Meta Graph 中恢复的图，其训练是从随机初始化的值开始的。训练中 Variable的实际值都保存在 .data和.index文件中，如果要从之前训练的状态继续恢复训练，就要从checkpoint 中 restore. tf.train.saver.save()在保存checkpoint的同时也会保存Meta Graph，但是在恢复图时，tf.train.saver.restore() 只恢复 Variable，如果要从MetaGraph恢复图，需要使用 import_meta_graph，当然我们也可以从模型的前向推断函数直接恢复图，这样我们只恢复Variable就好了。export_meta_graph/import_meta_graph 就是用来进行 Meta Graph 读写的API，下面展示从MetaGraph恢复图的方式：
with tf.Session() as sess: # load the meta graph saver = tf.train.import_meta_graph('./saved_model/model.meta') # get weights saver.restore(sess, tf.train.latest_checkpoint("./saved_model/"))
2.2：GraphDef简介
上文中，我们通过TensorBoard看到的计算图所表达的数据流其实与 python 代码中所表达的计算是对应的关系，但是在真实的 TensorFlow 运行中，Python 构建的Graph并不是启动一个session之后始终不变的东西。因为TensorFlow在运行时，真实的计算会被下放到多CPU上，或者 GPU 等异构设备，或者ARM等上进行高性能/能效的计算。实际上，TensorFlow而是首先将 python 代码所描绘的图转换（即“序列化”）成 Protocol Buffer（和 XML、JSON一样都是结构数据序列化的工具），再通过 C/C++/CUDA 运行 Protocol Buffer 所定义的图，该图叫GraphDef，GraphDef由许多叫做 NodeDef 的 Protocol Buffer 组成，在概念上 NodeDef 与 python代码中的操作相对应，保存网络的连接信息。通过tf.train.write_graph()/tf.Import_graph_def() 这一对api我们可以进行 GraphDef 的读写，它支持俩种不同的文件保存格式，下面展示文本格式的保存方法和结果：
import TensorFlow as tf # create variables a and b a = tf.get_variable("A", initializer=tf.constant(3)) b = tf.get_variable("B", initializer=tf.constant(5)) c = tf.add(a, b) saver = tf.train.Saver() with tf.Session() as sess: sess.run(tf.global_variables_initializer()) # 文本格式，as_text = True tf.train.write_graph(sess.graph_def, '.', 'model.pb', as_text = True)
如下所示是model.pb文件内的一个NodeDef的详情，其中包含name，op，input，attr等字段：
node {
name: "Add"
op: "Add"
input: "Const"
input: "Const_1"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "init"
op: "NoOp"
}
可以看到：GraphDef 中只有网络的连接信息（input字段表明该变量有俩个输入，分别是Const和Const_1），却没有任何 Variables，所以使用GraphDef 是不能够用来恢复训练的（没有权重）。
如果采用二进制格式（as_text = False）的方式来保存，生成文件会小得多，缺点就是它不易读。由此产生了俩种加载GraphDef的方式：
# as_text=False with tf.Session() as sess: with open('./model.pb', 'rb') as f: graph_def = tf.GraphDef() graph_def.ParseFromString(f.read()) # as_text=True from google.protobuf import text_format with tf.Session() as sess: # 不使用'rb'模式 with open('./model.pb', 'r') as f: graph_def = tf.GraphDef() text_format.Merge(f.read(), graph_def)
2.3：Frozen GraphDef简介
Frozen GraphDef，顾名思义，属于冻结（Frozen）后的 GraphDef 文件，这种文件格式不包含 Variables 节点。将 GraphDef 中所有 Variable 节点转换为常量（其值从 checkpoint 获取），就变为 Frozen GraphDef 格式。
下面展示如何利用GraphDef以及权重文件整合在一起获取 Frozen GraphDef：
import TensorFlow as tf from TensorFlow.python.tools import freeze_graph # network是你自己定义的模型 import network # 模型的checkpoint文件地址 ckpt_path = "./saved_model/" def freeze_graph_solution(): x = tf.placeholder(tf.float32, shape=[None, 224, 224, 3], name='input') # output是模型的输出 output = network(x) #设置输出类型以及输出的接口名字 flow = tf.cast(output, tf.int8, 'out') with tf.Session() as sess: #保存GraphDef tf.train.write_graph(sess.graph_def, '.', 'model.pb', as_text = True) #把图和参数结构一起,如果as_text为false，记得修改input_binary=True freeze_graph.freeze_graph( input_graph='./model.pb', input_saver='', input_binary=False, input_checkpoint=ckpt_path, output_node_names='out', restore_op_name='', filename_tensor_name='', output_graph='./frozen_model.pb', clear_devices=False, initializer_nodes='' )
其中用到了freeze_graph命令，这是一个非常有用的命令，后面还会接着出现。
除此之外，我们还可以将meta_graph_def以及权重文件整合在一起获取Frozen GraphDef：
# GraphDef 虽然不能保存 Variables，但是它可以保存constant with tf.Session() as sess: # load the meta graph and weights saver = tf.train.import_meta_graph('./model/model.meta') # get weights saver.restore(sess, tf.train.latest_checkpoint("./model/")) # 设置输出类型以及输出的接口名字 graph = convert_variables_to_constants(sess, sess.graph_def, ["out"]) tf.train.write_graph(graph, '.', 'frozen_model.pb', as_text = False)
3：面向部署的俩种模型结构
TensorFlow Serving是GOOGLE开源的一个服务系统，适用于部署机器学习模型，灵活、性能高、可用于生产环境。 它所依赖的模型结构是SavedModel ，SavedModel是TensorFlow 2.0 推荐的模型保存格式，一个 SavedModel 包含了一个完整的 TensorFlow program, 包含了 weights 以及 计算图，它不需要原本的模型代码就可以加载，很容易在 TFLite, TensorFlow.js, TensorFlow Serving, or TensorFlow Hub 上部署。一般github上提供的预训练模型文件也是这种结构的。下图展示了SavedModel在训练模型和部署模型中起到的重要作用。
Fig. 2. SavedModel在训练模型和部署模型中起到的重要作用
然而，有的时候我们对部署好的模型推理速度也有很高的要求，比如无人车驾驶场景中，如果使用一个经典的深度学习模型，很容易就跑到200毫秒的延时，那么这意味着，在实际驾驶过程中，你的车一秒钟只能看到5张图像，这当然是很危险的一件事。所以，对于实时响应比较高的任务，模型的加速就是很有必要的一件事情了。英伟达提供的Tensor-RT，就是一个高性能的深度学习推理（Inference）优化器，可以为深度学习应用提供低延迟、高吞吐率的部署推理。如下图所示，TensorRT现已能支持TensorFlow、Caffe、Mxnet、Pytorch等深度学习框架，将TensorRT和NVIDIA的GPU结合起来，能在几乎所有的框架中进行快速和高效的部署推理。但是除了caffe和TensorFlow，其他深度学习框架则需要先将模型转换为Open Neural Network Exchange（ONNX，开放神经网络交换）才可以。
Fig. 3. Tensor-RT对各大深度模型框架的支持
下面分别来介绍这俩种模型结构以及他们对应的部署方式。
3.1：SavedModel 方法与TensorFlow Serving
SavedModel 模型文件结构如下：
saved_model/
├── saved_model.pb
└── variables
├── variables.data-00000-of-00001
└── variables.index
顾名思义，variables保存所有变量，saved_model.pb保存的图即为GraphDef，包含模型结构等信息。
这种格式的模型有俩种保存方法，一种简单，一种虽然复杂但拥有更高的灵活性。
-
复杂方法：
# 保存 builder = tf.saved_model.builder.SavedModelBuilder('./saved_model') # x 为输入tensor inputs = {'input_x': tf.saved_model.utils.build_tensor_info(x)} # y 为最终需要的输出结果tensor outputs = {'output' : tf.saved_model.utils.build_tensor_info(y)} signature = tf.saved_model.signature_def_utils.build_signature_def(inputs, outputs, 'my_graph_tag') builder.add_meta_graph_and_variables(sess, ['test_saved_model'], {'my_graph_tag':signature}) builder.save()
-
简化的方法：
# 保存 tf.saved_model.simple_save(sess, model_path, inputs={'input_x': x_input}, outputs={'output': y_output})
使用tf.serving调用保存的模型：
image = cv2.imread(img_path) input_image_size = (513,513) # 数据预处理 X = image_util.general_preprocessing(image, 'tf',target_size=input_image_size) h, w = image.shape[:2] inputs = [X] input_names = ['input_x'] output_names = ['output'] outputs = request_tfserving(inputs=inputs, server_url='ip:port', model_name='deeplab_tf', signature_name='my_graph_tag', input_names=input_names, output_names=output_names)
细心的读者会发现，俩种方法里我们都重新定义了模型的输入输出的tensor的名字，而且在模型调用的时候名字保持了统一，那么如果我们的模型文件是下载好的预训练模型，我们并不知道模型的输入输出的tensor的名字怎么办？一种很简单的方式是利用saved_model_cli：
saved_model_cli show --all --dir ./saved_model # 输出saved_model的输入输出信息 signature_def['serving_default']: The given SavedModel SignatureDef contains the following input(s): inputs['image'] tensor_info: dtype: DT_FLOAT shape: (-1, 513, 513, 3) name: input_1:0 The given SavedModel SignatureDef contains the following output(s): outputs['result'] tensor_info: dtype: DT_FLOAT shape: (-1, 513, 513, 151) name: bilinear_upsampling_3/ResizeBilinear:0 Method name is: TensorFlow/serving/predict
可以看到，我们的模型输入的tensor名字是“input_1:0”，shape为 (-1, 513, 513, 3)，输出的tensor名字是“bilinear_upsampling_3/ResizeBilinear:0”，shape为(-1, 513, 513, 151)。
我们还可以深究一下什么是SignatureDef，它将输入输出tensor的信息都进行了封装，并且给他们一个自定义的别名，所以在构建模型的阶段，可以随便给tensor命名，只要在保存训练好的模型的时候，在SignatureDef中给出统一的别名即可，上文的示例中'my_graph_tag'就是我们所定义的SignatureDef的名字。
3.2：Frozen GraphDef与Tensor-RT
体会过了tf.serving带给我们服务部署上的便捷，接下来我们享受下Tensor-RT给我们模型带来的加速。
上文我们提到，一般github上提供的预训练模型文件是基于saved_model方法的，但是要实现模型加速需要获取Frozen GraphDef模型文件，它主要的用途是用于生产环境，或者发布产品等。所以我们要怎么整合saved_model格式的模型结构文件和权重文件呢？很简单，依旧可以使上文提到过的freeze_graph方法。
from TensorFlow.python.tools import freeze_graph from TensorFlow.python.saved_model import tag_constants # api freeze_graph.freeze_graph( input_graph=None, input_saver="", input_binary=False, input_checkpoint=None, output_node_names="out", restore_op_name='', filename_tensor_name='', output_graph='./frozen_model.pb', clear_devices=False, initializer_nodes='' input_saved_model_dir="./saved_model", saved_model_tags= tag_constants.SERVING )
freeze_graph也可以采用命令行的方式来执行，它和saved_model_cli一样，是安装好TensorFlow就可以使用的指令。
接下来展示如何利用ONNX文件来使用Tensor-RT进行模型推理加速的demo（详情可以参照https://developer.nvidia.com/zh-cn/tensorrt，里面也提供了很多主流模型的预训练模型，可以开箱即用）：
import TensorFlow as tf import tensorrt as trt import pycuda.driver as cuda import pycuda.autoinit import uff import image_util import common TRT_LOGGER = trt.Logger(trt.Logger.WARNING) # 构建引擎. def build_engine(onnx_file_path,engine_file_path): with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network, trt.OnnxParser(network, TRT_LOGGER) as parser: builder.max_workspace_size = 1 << 30 # 1GB builder.max_batch_size = 1 with open(onnx_file_path, 'rb') as model: parser.parse(model.read()) last_layer = network.get_layer(network.num_layers - 1) network.mark_output(last_layer.get_output(0)) # Build and return an engine. engine = builder.build_cuda_engine(network) with open(engine_file_path, "wb") as f: f.write(engine.serialize()) return engine image = cv2.imread(img_path) input_image_size = (513,513) image = image_util.general_preprocessing(image, 'tf',target_size=input_image_size) with build_engine(onnx_file_path,engine_file_path) as engine, engine.create_execution_context() as context: inputs, outputs, bindings, stream = common.allocate_buffers(engine) inputs[0].host = image trt_outputs = common.do_inference(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream)
4：结束语
