mathlf2015 / tfkreas2onnx

keras model trans to onnx

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

tf/keras模型线上推理加速方法:

方案1.tf2onnx转换tf/keras模型至onnx格式,用onnxruntime来加速在线推理

  • 示例:可以参照tfkeras2onnx.py文件,需要事先安装tf2onnx,如下:
    pip install tf2onnx
  • 需要首先将模型存储成pb文件,保存的文件名必须是xxxx/saved_model.pb,tf2onnx只支持这样的名称转换
  • 然后可以参照tfkeras2onnx.py文件,也可以使用如下方法直接转换:
    python -m tf2onnx.convert --saved-model "your_onnx_model_name/saved_model.pb" --output "save_root/your_model.onnx"
  • 加载(使用)保存模型方法,见bert_onnx_encoder.py

方案2.将tf/keras参数冻结,tf加载静态图推理

  • 参照freeze_pb.py文件,按照其中的示例来转换出模型的pb文件
  • 需要将模型用keras的model包一层,输入输出的长度固定死,然后再用freeze_keras_model2pb转换
  • 加载(使用)保存模型方法,见bert_freeze_encoder.py

二者对比压测实验

这里用到的bert是苏剑林开源的roberta small(4层),框架用的也是苏神的bert4keras,需要事先安装:

pip install bert4keras==0.10.8

压测实验效果:

用上面的两种方法将一个4层bert模型分别转换后压测实验,实验方法是用二者分别对相同的1w条query进行推理,查看压测时间,效果如下:

优化方法 P90/s P95/s P99/s P999/s
onnxruntime 0.00549 0.00672 0.00932 0.07780
freeze参数 0.00349 0.00373 0.00439 0.00897

结论:

  • 从压测来看,效果好像是freeze参数更好一些
  • 二者相同点是大部分推理耗时都在2.5ms以内
  • onnxruntime有个优点,有75%的推理耗时在3ms以内,而 freeze参数 则是只有60%的推理耗时在3ms以内
  • 个人推测onnxruntime效果差的原因可能是笔者用的cpu比较垃圾,导致前面的计算阻塞了,从而后面的计算堆积了起来。在真实场景下效果应该会比freeze参数稳定。

About

keras model trans to onnx


Languages

Language:Python 100.0%