之前成功将pytorch下的bert模型转换为onnx及tensorrt框架下的模型进行推理，特别在在tensor框架下，推理速度估算是提高了3-4倍（估算，跟显卡及设置的batch_size大小有关），参考的是这篇文章：https://blog.csdn.net/HUSTHY/article/details/118444462https://blog.csdn.net/HUSTHY/article/details/118444462 下面记录将keras模型转换为onnx模型。由于自己水平有限，以下内容仅为自己的理解并记录爬坑过程，肯定有更佳的方案。
1、将keras下的预训练模型加载后，进行推理时要用model.predict(x，y)的形式，想直接改成model(x,y)。自己首先进行了源码改写，出现错误：
RuntimeError: Detected a call to `Model.predict` inside a `tf.function`、`Model.predict is a high-level endpoint that manages its own `tf.function`、Please move the call to `Model.predict` outside of all enclosing `tf.function`s、Note that you can call a `Model` directly on `Tensor`s inside a `tf.function` like: `model(x)`.
后面查资料才知道需要用 tf.function 将方法包装起来起来，参考的做法：
import numpy as npimport tensorflow as tfdef build_model(): from tensorflow.keras.layers import Input, Multiply from tensorflow.keras.models import Model inputs = x = Input(shape=(4, 4, 3), name='image') x = Multiply()([x, x]) model = Model(inputs, x) return model# Create an array of all two's_input = np.full(shape=(1, 4, 4, 3), fill_value=2, dtype='float32')_model = build_model()print(_model.predict(_input))# correctly prints [[[[4、4、4.] ....# wrap model in a concrete function@tf.functiondef model_func(inputs): return _model(inputs)print(model_func(_input))
2、keras模型转换为onnx模型
参考代码
tensorflow-onnx/getting_started.py at master · onnx/tensorflow-onnx (github.com)https://github.com/onnx/tensorflow-onnx/blob/master/examples/getting_started.py 我的转换代码
output_model_path = 'keras_onnx_cpu.onnx'input_signature = [tf.TensorSpec([None, None], tf.float32), tf.TensorSpec([None, None], tf.float32)]onnx_model, _ = tf2onnx.convert.from_function(model, input_signature, opset=13, output_path=output_model_path)
3、加载onnx模型并进行推理
参考上面github中的加载代码会报错
sess = ort.InferenceSession(onnx_model.SerializeToString())
我用下面的代码
sess = ort.InferenceSession(output_model_path)res = sess.run(None, {'token_ids': token_ids.astype('float32'), 'segment_ids':segment_ids.astype('float32')})print(res[0])
此时如果报错
onnxruntime.capi.onnxruntime_pybind11_state.InvalidArgument: [ONNXRuntimeError] : 2 : INVALID_ARGUMENT : Failed to load model with error: Unknown model file format version.
在网上找的解决方案是：应该是onnx和onnxruntime的版本不匹配导致，可以安装onnx==1.8.0，onnxruntime==1.10.0，然后重新转换模型，再推理。我的换成onnx==1.8.0，onnxruntime==1.6.0后，加载成功。
4、关于GPU加速
在之前将prtorch模型转为onnx模型，用的是onnxruntime-gup(1.2.0)版，转换后的模型可用GPU加速。
本次转换需要 opset>=13，onnxruntime-gup的1.2.0版本太低，而我服务器的CUDA和cuDNN不支持高本版的onnxruntime-gup，所以只能用onnxruntime进行cpu推理。python onnx 快捷安装 onnxruntime 的 gpu 版本 如何使用_x1131230123的博客-CSDN博客_onnx python 安装https://blog.csdn.net/x1131230123/article/details/120422132 更新内容：在高版本GUDA（11.1）环境上安装了onnx==1.10.1，onnxruntime-gpu==1.8.1，用GPU加速onnx推理，发现速度提升更多。
5、推理速度比较
用自己的数据在服务器(显存)环境下跑，cpu下onnx模型的速度比gpu下原keras模型的速度快2-3倍，gpu下onnx模型的速度比gpu下原keras模型的速度快6倍左右。
6、未来可能的做法
如果后面有需要，进一步将onnx模型转换为tensortr模型，应该能再次提高推理速度。