Tensorrt的使用：pytorch-＞onnx-＞trt

nuo112

已于 2022-03-29 10:24:02 修改

阅读量2.3k

点赞数 1

CC 4.0 BY-SA版权

分类专栏：软件安装文章标签： ubuntu python 深度学习 pytorch

于 2022-03-29 10:15:58 首次发布

本文链接：https://blog.youkuaiyun.com/nuo112/article/details/123813103

软件安装专栏收录该内容

7 篇文章

订阅专栏

本文档详细介绍了如何在特定环境下将PyTorch模型转换为ONNX格式，并进一步转换为TensorRT（TRT）格式的过程。具体步骤包括设置Ubuntu18.04系统、CUDA10.2、cuDNN8.0.2等软件环境，解决动态batch问题，以及使用onnx2trt工具完成模型转换。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

0. 电脑环境(此环境可供参考，因为本人在配置时，尝试不少版本间的配合使用！！！难受！！！)

tensorrt的安装推荐

系统：Ubuntu 18.04
cuda: 10.2
cudnn: 8.0.2(推荐配置8.2.1)
torch: 1.10.0
torchvision: 0.11.2
Tensorrt: 8.2.3.0(推荐此版本，在我使用中，其他版本的容易有兼容性问题)

1. pytorch转onnx参考如下博客进行转换：

Pytorch分类模型转onnx以及onnx模型推理 - 知乎（一）Pytorch分类模型转onnx 参考：PyTorch之保存加载模型PyTorch学习：加载模型和参数_lscelory的博客-优快云博客_pytorch 加载模型实验环境：Pytorch1.4 + Ubuntu16.04.5 1.Pytorch之保存加载模型1.1 当提到保存…https://zhuanlan.zhihu.com/p/159379768 按照上面的方法使用，在进行onnx转trt时会报动态batch的错误，因此，需要改成静态batch，把torch.onnx.export()部分改写为：

torch.onnx.export(model_yolov3,
                    x,
                    export_onnx_file,
                    opset_version=11,
                    export_params=True,
                    input_names=["input"],	# 输入名
                    output_names=["output"]	# 输出名
                    )

2. 生成onnx文件之后，需要转成trt文件进行模型的推理。

onnx转trt参考如下博客进行转换,使用onnx2trt工具进行转换：

TensorRT教程4：使用onnx2trt工具转engine_米斯特龙_ZXL的博客-优快云博客_onnx2trt使用onnx2trt工具转engineonnx-tensorrt 地址：https://gitee.com/Mr_xiaolong/onnx-tensorrt1、安装onnx2trt注意我用的trt7.2.3.4版本#克隆master分支到本地，因为默认master分支是7.2.3.4git clone https://gitee.com/Mr_xiaolong/onnx-tensorrt.gitgit submodule init#子模块初始化git submodule update #更新https://blog.youkuaiyun.com/weixin_41562691/article/details/118277630 注：我使用的是cuda10.2，有2个补丁需要下载安装，如果未安装，会报下面的错误。

[2021-07-26 07:16:07   ERROR] 2: [ltWrapper.cpp::setupHeuristic::327] Error Code 2: Internal Error (Assertion cublasStatus == CUBLAS_STATUS_SUCCESS failed.)
terminate called after throwing an instance of 'std::runtime_error'
  what():  Failed to create object
已放弃 (核心已转储)

安装方式为

1. 下载补丁1、2

2. 运行：sudo sh cuda_10.2.1_linux.run 进行安装。安装时，去掉CUDA Toolkit 10.2选项，因为已经安装。