摘要
基于tensorflow训练的模型一般被保存为ckpt形式的文件,随着当前深度学习模型网络越来越大,对应模型也会非常大。当对外提供服务的时候,如果采用ckpt的形式,服务进程被调起来非常困难,且推理服务一般速度也较慢(会达到100ms以上一次推理)。所以,使用固化后的模型进行推理服务非常必要,模型大小得到压缩,服务时耗也得到大幅度降低。本文介绍几种常见的模型固化方法。
1 session模式的模型
先行训练好check point形式的模型文件,之所以这里区分session模式,因为这种模式使用者比较容易设置name_scope以及Tensor的命名,
特别是输入、输出节点的名称,这样对于后面的固化非常重要。代码如下:
#!/usr/bin/env python
# coding=utf-8
from __future__ import print_function
import os,sys,time
import numpy as np
import tensorflow as tf
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "7"
import tensorflow.contrib.slim as slim
from tensorflow.python.framework import graph_util
#此文件可以把ckpt模型转为pb模型
def freeze_graph(input_checkpoint,output_graph):
# checkpoint = tf.train.get_checkpoint_state(model_folder) #
# input_checkpoint = checkpoint.model_checkpoint_path #
output_node_names = "score_student/output_student"
saver = tf.train.import_meta_graph(input_checkpoint + '.meta', clear_devices=True)
graph = tf.get_default_graph()#
input_graph_def = graph.as_graph_def()#
with tf.Session() as sess:
saver.restore(sess, input_checkpoint)
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(","),
variable_names_whitelist=None,variable_names_blacklist=None)#
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))
#
input_checkpoint='/data/.../best_validation'
out_pb_path='/data/.../pbmodel/frozen_model_for_best_validation20190821.pb'
freeze_graph(input_checkpoint, out_pb_path)
2 estimator形式的模型
包括TPUEstimator形式,最常见的就是
BERT模型[1]了,你很难去找里面的各种tensor的变量(本人亲自尝试使用tensorboard进行查找)。【
当然使用一次迁移学习转换一下形式也是可以解决这个问题的】,这里提供另外一种简单的解决方法,代码如下:【相关表达方式还有好几种】
首先使用正常模型训练一个ckpt形式的模型,然后进行转化pb格式固化。在bert代码中注释掉原来的train步骤
#estimator.train(input_fn=train_input_fn, max_steps=num_train_steps)
#定义里面的输入数据名称及shape:
feature_spec = {
"input_ids": tf.FixedLenFeature([FLAGS.max_seq_length], tf.int64),
"input_mask": tf.FixedLenFeature([FLAGS.max_seq_length], tf.int64),
"segment_ids": tf.FixedLenFeature([FLAGS.max_seq_length], tf.int64),
"start_ps": tf.FixedLenFeature([], tf.int64),#@debuluoyi
"end_ps": tf.FixedLenFeature([], tf.int64),#@debuluoyi
"label_ids": tf.FixedLenFeature([], tf.int64), }
#固化转换
receiver_fn=tf.estimator.export.build_parsing_serving_input_receiver_fn(feature_spec)
estimator._export_to_tpu = False
print("begin to export_savedmodel el model.pb") estimator.export_savedmodel('/data/.../pbmodel/bert_el_model20191023_2_pbmodel',serving_input_receiver_fn =receiver_fn)
3 小结
模型固化对于对外提供服务很重要,另外还有一个很重要的点就是接下来进行压缩,压缩很多基于pb固化模型进行也有多种方便,
由此体现模型固化在实际算法工程中是非常重要且常见的技术点,一般能够将模型大小压缩数倍,服务调起更加轻便,基于GPU的推理时间也能节约数倍。参考文献
[1] https://github.com/google-research/bert;
google-research/bertgithub.com
[2] 本人GitHub:
debuluoyi - Overviewgithub.com