摘要
基于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/bert?github.com
[2] 本人GitHub：
debuluoyi - Overview?github.com
)
