(ubuntu 18.04 cuda 10.2 pytorch onnx→tensorrt) 使用tensorrt加载并运行onnx模型

最新推荐文章于 2025-04-10 16:58:12 发布

kiyoxi

最新推荐文章于 2025-04-10 16:58:12 发布

阅读量952

点赞数 1

分类专栏：神经网络文章标签：深度学习 cuda pytorch

本文链接：https://blog.youkuaiyun.com/heiha1232/article/details/115730431

版权

本文介绍了如何在Ubuntu 18.04环境下，使用CUDA 10.2和TensorRT 7.0安装TensorRT，并将PyTorch模型转换为ONNX格式，然后利用TensorRT加载并运行ONNX模型进行预测。主要流程包括TensorRT和pycuda的安装，PyTorch到ONNX的转换，以及TensorRT调用ONNX模型进行预测。需要注意的是，TensorRT对pad层的要求以及7.0.0.11版本不支持batchnorm层。

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

2021.4

ubuntu 18.04 cuda 10.2 pytorch onnx→tensorrt使用tensorrt加载并运行onnx模型

0. TensorRT介绍
1. 安装tensorrt、pycuda
- 1.1 安装tensorrt
- 1.2 安装pycuda
2. 将pytorch模型转化为onnx模型
3. tensorrt调用onnx模型进行预测
4. 注意事项

0. TensorRT介绍

https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html

TensorRT的核心是一个C ++库，可促进对NVIDIA图形处理单元（GPU）进行高性能推断。它旨在与TensorFlow，Caffe，PyTorch，MXNet等深度学习框架互补地工作。它专门致力于在GPU上快速有效地运行已经受过训练的网络，以生成结果（包括检测，回归或推断任务等）。
TensorRT根据指定的精度（FP32，FP16或INT8）组合各层，优化内核选择，通过归一化、转换为更优化的矩阵数学运算等方式，改善延迟，吞吐量和效率。
主要流程：

1. 安装tensorrt、pycuda

1.1 安装tensorrt

环境：ubuntu18.04, 2080Ti，安装cuda10.2+cudnn 7.6.5，且在~/.bashrc中添加cuda、cudnn所在位置：
tensorrt下载链接：https://developer.nvidia.com/zh-cn/tensorrt

需要选择tensorrt 7.0.0.11（适用于Ubuntu18.04、cuda10.2、cudnn7.6，其他更高的tensorrt版本要求cudnn8，因此此处选择tensorrt 7.0.0.11）
- 下载之后解压：
```
tar xzvf TensorRT-7.0.0.11.Ubuntu-18.04.x86_64-gnu.cuda-10.2.cudnn7.6.tar.gz
```
- 解压后得到“TensorRT-7.0.0.11”文件，包含以下目录：

安装指导：打开doc/pdf/TensorRT-Installation-Guide.pdf

由于下载的是tar安装文件，因此根据文件中4.3 tar file installation章节安装即可

主要安装步骤如下：

export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:~/TensorRT-7.0.0.11/lib   #TensorRT文件所在目录

cd TensorRT-7.0.0.11/python
pip3 install tensorrt-*-cp3x-none-linux_x86_64.whl # 根据python版本选择whl文件

cd TensorRT-7.0.0.11/graphsurgeon
pip3 install graphsurgeon-0.4.1-py2.py3-none-any.whl

验证安装：进入samples目录，该目录下有很多c++、python的tensorrt使用示例，可以详细学习一下
- 进入samples/sampleMNIST目录
```
make clean   # 可能需要sudo
make 

cd ../../bin
./sample_mnist
```
- 还需要准备数据：进入data/minist目录下，运行
```
python download_pgms.py
```
  得到0.pgm~9.pgm，注意download_pgms.py里面的mnist下载链接可能无效，此时可以找一个有效的链接进行替换
- 得到结果

1.2 安装pycuda

pip install pycuda

2. 将pytorch模型转化为onnx模型

tensorrt无法直接读取pytorch模型，因此需要先转化为onnx

以alexnet为例，主要代码：

import torch
import torchvision

def get_model():
    model = torchvision.models.alexnet(pretrained=True).cuda()
    resnet_model = model.eval()
    return model

def get_onnx(model, onnx_save_path, example_tensor):
    example_tensor = example_tensor.cuda()
    _ = torch.onnx.export(model, 
                        example_tensor, 
                        onnx_save_path,
                        verbose=True,
                        training=False,
                        do_constant_folding=False,
                        input_names=['input'],
                        output_names=['output'])

if __name__ == '__main__':
    model = get_model()
    onnx_save_path = "alexnet.onnx"
    example_tensor = torch.randn

最低0.47元/天解锁文章