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:模型概览
Graph)的方式来表达,TensorFlow官网有这样的一句介绍:“A computational graphnode也被称为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()
MetaGraph
恢复图,需要使用 import_meta_graph
,当然我们也可以从模型的前向推断函数直接恢复图,这样我们只恢复Variable就好了。export_meta_graph/import_meta_graph
Meta Graph
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/"))
View Code
2.2:GraphDef简介
XML
、JSON
一样都是结构数据序列化
的工具),再通过 C/C++/CUDA 运行 Protocol Buffer 所定义的图,该图叫GraphDef,GraphDef由许多叫做 NodeDeftf.train.write_graph()/tf.Import_graph_def()
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)
View Code
如下所示是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"
}
View Code
可以看到: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)
View Code
2.3:Frozen GraphDef简介
Frozen GraphDef
,顾名思义,属于冻结(Frozen)后的 GraphDef 文件,这种文件格式不包含 Variables 节点。将 GraphDef 中所有 Variable 节点转换为常量(其值从 checkpoint 获取),就变为 Frozen 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=''
)
View Code
其中用到了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)
View Code
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()
View Code
- 简化的方法:
# 保存
tf.saved_model.simple_save(sess, model_path, inputs={'input_x': x_input}, outputs={'output': y_output})
View Code
使用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)
View Code
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
View Code
可以看到,我们的模型输入的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
)
View Code
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)
View Code
4:结束语
研究TensorFlow代码的时候,最苦恼的就是明明只想实现一个简单的功能,为什么会有好多种实现方式,好不容易搞明白了这种方法,又看到了大佬采用另外一种更优雅的实现方式,想去研究的时候,发现各种概念错综复杂,之前花了很大功夫把模型文件的相关概念理了一通,并简单介绍了tf.serving和Tensor-RT的入门级使用方法,仅做抛砖引玉之用,希望能对大家有所帮助,少踩一点坑,可能会有遗漏和不足,敬请指出。
源码地址:https://github.com/LeiyuanMa/TensorFlow_1.x_model