摘要

基于tensorflow訓練的模型一般被保存為ckpt形式的文件，隨著當前深度學習模型網(wǎng)絡越來越大，對應模型也會非常大。當對外提供服務的時候，如果采用ckpt的形式，服務進程被調(diào)起來非常困難，且推理服務一般速度也較慢(會達到100ms以上一次推理)。所以，使用固化后的模型進行推理服務非常必要，模型大小得到壓縮，服務時耗也得到大幅度降低。本文介紹幾種常見的模型固化方法。

1 session模式的模型

先行訓練好check point形式的模型文件，之所以這里區(qū)分session模式，因為這種模式使用者比較容易設置name_scope以及Tensor的命名，特別是輸入、輸出節(jié)點的名稱，這樣對于后面的固化非常重要。代碼如下：

#!/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模型轉(zhuǎn)為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進行查找）。【當然使用一次遷移學習轉(zhuǎn)換一下形式也是可以解決這個問題的】，這里提供另外一種簡單的解決方法，代碼如下：【相關(guān)表達方式還有好幾種】

首先使用正常模型訓練一個ckpt形式的模型，然后進行轉(zhuǎn)化pb格式固化。

在bert代碼中注釋掉原來的train步驟

#estimator.train(input_fn=train_input_fn, max_steps=num_train_steps)

#定義里面的輸入數(shù)據(jù)名稱及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), }

#固化轉(zhuǎn)換

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 小結(jié)

模型固化對于對外提供服務很重要，另外還有一個很重要的點就是接下來進行壓縮，壓縮很多基于pb固化模型進行也有多種方便，由此體現(xiàn)模型固化在實際算法工程中是非常重要且常見的技術(shù)點，一般能夠?qū)⒛Ｐ痛笮嚎s數(shù)倍，服務調(diào)起更加輕便，基于GPU的推理時間也能節(jié)約數(shù)倍。

參考文獻

[1] https://github.com/google-research/bert；

google-research/bert?github.com

[2] 本人GitHub：

debuluoyi - Overview?github.com

總結(jié)

以上是生活随笔為你收集整理的tensorboard ckpt pb 模型的输出节点_算法工程化系列——模型固化的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯，歡迎將生活随笔推薦給好友。