ONNX-TensorRT简介
ONNX-TensorRT是一个开源项目,主要功能是将ONNX (Open Neural Network Exchange)格式的深度学习模型解析并执行于NVIDIA的TensorRT推理引擎上。它作为ONNX的TensorRT后端实现,可以将ONNX模型无缝转换为TensorRT引擎,充分利用TensorRT的优化功能来加速模型推理。
ONNX-TensorRT的主要特点包括:
- 支持大多数常见的ONNX算子,可以解析各种类型的神经网络模型
- 利用TensorRT的优化技术,如层融合、精度校准等,显著提升模型推理性能
- 支持动态shape和全维度张量,增强了模型的灵活性
- 提供C++和Python两种编程接口,方便集成到不同的应用中
- 持续更新以支持最新版本的ONNX和TensorRT
支持的TensorRT版本
当前分支开发针对的是最新的TensorRT 10.2版本,支持全维度和动态shape。对于之前的TensorRT版本,可以参考相应的历史分支。
支持的ONNX算子
ONNX-TensorRT支持大部分常用的ONNX算子,具体支持列表可以参考项目文档中的算子支持矩阵。随着版本迭代,支持的算子还在不断扩充中。
安装与构建
依赖
- Protobuf >= 3.0.x
- TensorRT 10.2
- TensorRT 10.2 开源库
构建步骤
- 克隆ONNX-TensorRT代码仓库
- 创建build目录并进入
- 运行cmake配置,指定TensorRT路径
- 执行make编译
cd onnx-tensorrt
mkdir build && cd build
cmake .. -DTENSORRT_ROOT=<path_to_trt> && make -j
export LD_LIBRARY_PATH=$PWD:$LD_LIBRARY_PATH
注意需要正确设置CUDA路径,默认是/usr/local/cuda。
使用方法
C++接口
可以使用trtexec工具快速测试ONNX模型是否可以成功解析为TensorRT引擎:
trtexec --onnx=model.onnx
Python接口
可以使用polygraphy工具测试ONNX模型:
polygraphy run model.onnx --trt
ONNX-TensorRT的Python模块封装在.whl文件中,可以直接pip安装使用。
性能优化
ONNX-TensorRT提供了多种优化选项来提升模型性能,如:
- 支持FP16和INT8精度
- 层融合优化
- 动态shape支持
- 针对特定算子的优化实现
用户可以根据具体需求选择合适的优化策略。
总结
ONNX-TensorRT为将ONNX模型部署到TensorRT上提供了便捷的解决方案。通过将ONNX的灵活性与TensorRT的高性能相结合,可以显著提升深度学习模型在NVIDIA GPU上的推理效率。该项目仍在持续更新中,未来会支持更多ONNX算子和TensorRT的新特性。
