keras保存模型_onnx+tensorrt部署keras模型

由于项目需要，最近捣鼓了一波如何让用tensorrt部署训练好的模型来达到更快的推理速度，期间花费了大量的时间在知乎和各种网页上去搜索别人的方案，但始终没有找到我想要的条理相对清晰的记录贴（也许只是我没找到），因此作为记录，也希望能帮到别人，弯路嘛，少走一条是一条。

由于之前项目一直是用tensorflow和keras，因此本文所讲述的是如何将模型从keras上迁移到tensorrt中，采用的方案是keras保存的hdf5模型先转为onnx模型，再将onnx模型转为tensorrt模型。采取这种方案主要有两点考虑，一方面是在tensorrt的官网只看到tensorflow转为tensorrt的教程，没有将keras模型直接部署的教程。另一方面是考虑到以后可能转去pytorch阵营（嘴上说着转已经说了半年，tf的api实在看的我眼花缭乱），到时候只需要再熟悉一下pytorch转onnx即可。废话不多说，下面是我的部署记录。

首先tensorrt是老黄为自家英伟达GPU所专门打造的推理引擎，TPU、CPU、FPGA等等请去看看TVM，本文是基于tensorrt已经安装好的基础上进行模型部署，没有安装的小伙伴请去百度，官网上的英文教程很详细，配合网友的中文安装记录很快就能配好。

Keras To Onnx

这一步主要是将keras保存的hdf5模型转为onnx后缀的模型，主要代码如下：

import keras2onnx
import onnx
from keras.models import load_model
model = load_model('testonnx.hdf5')  
onnx_model = keras2onnx.convert_keras(model, model.name)
temp_model_file = 'result.onnx'
onnx.save_model(onnx_model, temp_model_file)

代码很简单，分为三步，分别是加载keras模型，keras模型中转为onnx模型，保存onnx模型。需要注意的是加载的hdf5文件中需要包含网络结构和参数，仅仅包含参数是不行的。

此外，在我部署过程中，报了一些错误，主要是一些自定义的layer和自定义的loss无法被识别，如何转换自定义的函数我暂时还没有深入研究，后续会补上，采取了一些“曲线救国”的法子，把自定义的layer尽量用keras自带的代替，效果差点就差点吧。在训练时采用自定义的loss，最后用keras自带的loss重新编译加载再保存就不会再报错了。

如果不出错的话，生成onnx结尾的文件表示转换完成，接下来是如何将onnx模型转换为tensorrt模型。

Onnx To Tensorrt

在安装tensorrt时