背景：
用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
- *使用如下方法，可将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]
[,
]
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）。
如果有用，请留言并让我知道！如果有更好的建议或技术上的不足，请更正！谢谢!
[En]
If it’s useful, please leave a message and let me know! If there are any better suggestions or technical deficiencies, please correct them! Thank you!
Original: https://blog.csdn.net/b285795298/article/details/124296547
Author: 半九拾
Title: [NLP]keras模型部署加速（ONNX Speed Keras_Model Inference）
原创文章受到原创版权保护。转载请注明出处：https://www.johngo689.com/497508/
转载文章受原作者版权保护。转载请注明原作者出处！