[NLP]keras模型部署加速（ONNX Speed Keras_Model Inference）

目录

使用如下方法，可将Keras_model写成pb文件

tf2onnx 将pbfile转成onnx

---

背景：

用keras 训练 NER 模型，模型架构为：Roberta_CRF。

整个模型保存下来，1.2G 左右。

欲部署在cpu环境下，需要对模型进行压缩和加速。没有找到bert的蒸馏后的中文预训练模型。

整体方案是：keras_weights -> tf_pb_model(freeze graph) -> tf2onnx

ps:直接使用keras2onnx，因为transformer_op名称的问题，报错，并找不到解决方案。探索再三，使用上述方案。

环境：

tensorflow==1.12.0
python==3.6
keras==2.2.4
h5py==2.10.0
onnx==1.11.0
tf2onnx==1.10.0
onnxruntime==1.10.0

keras2onnx 报错为：

keras2onnx.convert_keras Transformer-11-FeedForward-Add_1/All:0 already processed

Tensor Transformer-11-MultiHeadSelfAttention-Add/All:0 already processed · Issue #609 · onnx/keras-onnx · GitHubhttps://github.com/onnx/keras-onnx/issues/609

• 使用如下方法，可将Keras_model写成pb文件

def freeze_session(session, keep_var_names=None, output_names=None, clear_devices=True):
    """
    Freezes the state of a session into a pruned computation graph.
    """
    from tensorflow.python.framework.graph_util import convert_variables_to_constants
    graph = session.graph
    with graph.as_default():
        freeze_var_names = list(set(v.op.name for v in tf.global_variables()).difference(keep_var_names or []))
        output_names = output_names or []
        output_names += [v.op.name for v in tf.global_variables()]
        input_graph_def = graph.as_graph_def()
        if clear_devices:
            for node in input_graph_def.node:
                node.device = ""
        #frozen_graph = tf.graph_util.convert_variables_to_constants(
        #    session, input_graph_def, output_names, freeze_var_names)
        frozen_graph = convert_variables_to_constants(session, input_graph_def, output_names, freeze_var_names)
        return frozen_graph


from keras import backend as K
import tensorflow as tf

'''
1、构造出graph
    from keras.models import Model
    model = Model(model.input, output)

2、加载权重(使用model.save_weights() 保存的训练好的模型权重)
    model.load_weights(self.saved_model_weights_dir)

3、写成pb file

'''

frozen_graph = freeze_session(K.get_session(),output_names=[out.op.name for out in model.outputs])
tf.train.write_graph(frozen_graph,'folder_to_save','saved_model.pb', as_text=False)

可以测试pb文件能够正常推理

import tensorflow as tf
from tensorflow.python.platform import gfile

with tf.Session() as sess:
    with gfile.FastGFile("pb_file_dir", 'rb') as f:
        graph_def = tf.GraphDef()
        # Parses a serialized binary message into the current message.
        graph_def.ParseFromString(f.read())
        sess.graph.as_default()
        # Import a serialized TensorFlow `GraphDef` protocol buffer
        # and place into the current default `Graph`.
        gin = tf.import_graph_def(graph_def)

        tensor_input_1 = sess.graph.get_tensor_by_name("import/Input-Token:0")
        tensor_input_2 = sess.graph.get_tensor_by_name("import/Input-Segment:0")
        tensor_output = sess.graph.get_tensor_by_name("import/conditional_random_field_1/add:0")
        model_predict = sess.run(tensor_output,{tensor_input_1:token_ids, tensor_input_2:segment_ids})

如果你不知道你的输入和输出op的名称,推荐如下方法，注意get_tensor_by_name 中，名称后加上":0"。

gg = sess.graph.get_operations()
gg[:2]
# [<tf.Operation 'import/Input-Token' type=Placeholder>,
# <tf.Operation 'import/Input-Segment' type=Placeholder>]
gg[-5:]
...

• tf2onnx 将pbfile转成onnx

安装tf2onnx

$ git clone https://github.com/onnx/tensorflow-onnx.git
$ cd tensorflow-onnx
$ python setup.py install

转onnx文件：

python -m tf2onnx.convert --graphdef  dir_to_pb_file/saved_model.pb --opset 13 --output ner_model.onnx --output model.onnx --inputs Input-Token:0,Input-Segment:0 --outputs conditional_random_field_1/add:0

要使用--graphdef 方式，要指明输入输出，要注意，这里的输入输出的名称前面的"import/" 就不需要了，否则报错graph中找不到这些层。

可以测试当前的onnx 文件是可以正常推理的：

import onnxruntime
session = onnxruntime.InferenceSession("ner_model.onnx")
session.get_modelmeta()
first_input_name = session.get_inputs()[0].name
second_input_name = session.get_inputs()[1].name
output_name = session.get_outputs()[0].name
results = session.run([output_name], 
                      {
                      first_input_name:token_ids.astype(np.float32), 
                    second_input_name:segment_ids.astype(np.float32)})

注意这个output_name 要用‘[]’ 包起来。 输入的numpy转成np.float32。

模型转onnx模型后，总体400M，推理的时候不需要指定bert_pretrain_model,速度也有相应提升。onnx跨平台友好，部署环境里可以少装几个包（如tensorflow）。

为什么onnx快？Make your model faster on CPU using ONNX - UbiOps

单GPU推理，如何优化？NVIDIA TensorRT | NVIDIA DeveloperNVIDIA TensorRT What is NVIDIA TensorRT Features Key Products Performance Framework Integrations Success Stories Introductory Resources NVIDIA TensorRT NVIDIA® TensorRT™, an SDK for high-performance deep learning inference, includes a deep learning inference optimizer and runtime that delivers low latency and high throughput for inference applications.https://developer.nvidia.com/tensorrt

有用的话，麻烦留言告诉我！有更好的建议或是技术上的不足，敬请指正！谢谢！