tensorboard ckpt pb 模型的输出节点_算法工程化系列—

tensorboard ckpt pb 模型的输出节点_算法工程化系列——模型固化

摘要

基于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)