TensorRT（C++）部署 Pytorch模型

最新推荐文章于 2026-07-01 15:57:42 发布

原创

最新推荐文章于 2026-07-01 15:57:42 发布 · 5.3k 阅读

本文详细介绍了如何在Ubuntu 20.04系统中安装TensorRT 8.4.1，配合CUDA 11.6和cudnn 8.4.1，用于yolov5模型的部署，包括环境设置、yolov5模型转换、编译并封装算法接口。特别提到了不同版本匹配的重要性。

众所周知，python训练pytorch模型得到.pt模型。但在实际项目应用中，特别是嵌入式端部署时，受限于语言、硬件算力等因素，往往需要优化部署，而tensorRT是最常用的一种方式。本文以yolov5的部署为例，说明模型部署在x86架构上的电脑端的流程。（部署在Arm架构的嵌入式端的流程类似）。

一、环境安装

1. 安装tensorRT

根据自己的系统Ubuntu版本、CPU架构、CUDA版本来选择tensorrt的下载包，连接如下

https://developer.nvidia.com/nvidia-tensorrt-8x-download

注意如果是自己的cuda不是用.deb安装的，最好是下载tensorrt的.tar包进行安装。笔者因为事先配置过pytorch开发环境，安装过cuda(.run)方式，所以使用.deb安装tensorrt时一直报错且未解决，后来使用.tar安装就好了。下载tensorrt的时候要注意版本对应关系，这里的版本对应关系包括cuda和cudnn的版本。如tensorrt使用的TensorRT-8.4.1.5.Linux.x86_64-gnu.cuda-11.6.cudnn8.4.tar.gz，则cuda要使用11.6，cudnn要使用8.4.1，否则编译tensorrt时会报错。

下载后文件为 TensorRT-8.4.1.5.Linux.x86_64-gnu.cuda-11.6.cudnn8.4.tar.gz

执行下列命令

# 解压
tar zxf TensorRT-8.4.1.5.Linux.x86_64-gnu.cuda-11.6.cudnn8.4.tar.gz
# 换个目录
cp TensorRT-8.4.1.5 /opt

编辑.bashrc文件把tensorrt的库文件添加到环境变量

# 开大bashrc文件
sudo gedit ~./bashrc
# 在文件末尾添加下列一行
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/opt/TensorRT-8.4.1.5/lib

保存关闭后，完成配置更新环境

source ~/.bashrc

复制tensorRT目录下 lib、include文件夹到系统文件夹（或者将这两个文件夹路径添加到.bashrc文件中）

sudo cp -r ./lib/* /usr/lib
sudo cp -r ./include/* /user/include

安装完成后，测试是否安装成功，这里使用 /opt/TensorRT-8.4.1.5/samples/sampleMNIST来测试

cd /opt/TensorRT-8.4.1.5/samples/sampleMNIST
sudo make

编译之后会在 /opt/TensorRT-8.4.1.5/bin生成对应的可执行文件，此时直接运行sample_mnist

sudo ./sample_mnist

得到下列结果即为成功

2. 安装opencv

具体参考我的另一篇博客

ubuntu编译安装opencv4.5.1(C++)_专业女神杀手的博客-CSDN博客

二、tensorrt的yolov5部署

tensorrt是一个推理引擎架构，会将pytorch用到的网络模块，如卷积，池化等用tensorrt进行重写。pytorch模型转换为.engine后就可以进行推理。

在github上下载tensorrt，注意tensorrt版本和yolov5版本的匹配关系，如tensorrt5.0对应yolov5的 v5.0

标签

#pytorch #人工智能 #c++

最低0.47元/天解锁文章