从Keras模型转换成tensorrt的方法

最新推荐文章于 2024-08-16 08:28:10 发布

jdkholmes

最新推荐文章于 2024-08-16 08:28:10 发布

阅读量3k

点赞数 2

文章标签： python tensorflow 人工智能深度学习

本文链接：https://blog.csdn.net/jdkholmes/article/details/108406470

版权

本文介绍了将Keras模型转换为TensorRT engine的详细流程，包括使用tf2onnx将savedmodel转换为ONNX模型，再通过onnx2trt生成engine。涉及的环境包括TensorFlow 1.14.0、CUDA 10.0和TensorRT 7.0.0.11等。在转换过程中遇到的问题，如模型维度不匹配和dynamic shape的处理，以及解决方案进行了说明。

摘要由CSDN通过智能技术生成

1 流程
对于目前找到的最有效的方式是将Keras模型保存为savedmodel格式，也就是说一个只含有计算图的.pb，一个variables文件和一个json合并的一个文件夹。tf-trt加速的话会还给你一个优化过后的.pb和variable，并不是我们希望使用的engine形式。
最稳妥的流程为：savedmodel–>onnx–>engine。
下载tf2onnx，使用如下命令：
python3 -m tf2onnx.convert --saved-model /lsj/checkpoint_germany_savedmodel/ --opset 7 --inputs input_1:0[1,224,224,3] --output /lsj/germany_model.onnx
具体参数GitHub有说明。
注意：一定要在root目录下运行此命令，否则会报错关于__path__。
最后使用onnx-trt工具，命令如下：
onnx2trt /lsj/germany_model.onnx -o /lsj/germany.engine
即可成功生成engine
2环境
TensorFlow 1.14.0
CUDA 10.0
TensorRT 7.0.0.11
cmake 3.18.2
protobuf 3.12.2