pytorch模型tensorrt加速之-pth转onnx转trt，在推理trt模型测试模型速度

最新推荐文章于 2024-09-24 14:48:26 发布

pyswt

最新推荐文章于 2024-09-24 14:48:26 发布

阅读量7k

点赞数 16

分类专栏：深度学习文章标签： python pytorch tensorrt cuda

本文链接：https://blog.csdn.net/python3_i_know/article/details/116262690

版权

本文介绍了如何将PyTorch模型通过TensorRT进行加速，详细讲解了从.pth模型转换为.onnx，再转为.trt模型的过程。在Ubuntu 18.4环境下，使用Python 3.8、PyTorch 1.8+cu11和TensorRT 7.2.3，通过torch2trt库完成模型转换。转换后的TRT模型在推理时速度比原始.pth模型快了4-5倍。

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

首先需要安装两个必要的包tensorrt和torch2trt，tensorrt的话要在官网下载tar压缩包，建议下载tar包安装，官网，我下载的是7.2.3版本，torch2trt可以在GitHub上克隆项目。
我的环境(tensorrt目前好像Linux上可以用，我在win10上没有搞成功，tar包里面没有对应的whl文件)：
ubuntu18.4
python3.8
pytorch1.8+cu11
tensorrt7.2.3
pycuda2021

解压安装tensorrt

tar -zxvf TensorRT-7.2.3.4.Ubuntu-18.04.x86_64-gnu.cuda-11.1.cudnn8.1.tar.gz
cd TensorRT-7.2.3.4
# 这个目录会有四个文件，python，uff，graphsurgeon，onnx_graphsurgeon，分别进入目录安装对应的whl文件
pip install xxxxx.whl

# 配置环境变量
vim /etc/profile
# 最后面加上如下命令，路径需看自己的安装路径，注意这个lib文件夹可能是个软连接，我当时因为这个报错了，要lib的正真路径
export TRT_RELEASE=`pwd`/TensorRT-7.2.3.4
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$TRT_RELEASE/lib

# 使文件立刻生效
source /etc/profile

安装torch2trt

git clone https://github.com/NVIDIA-AI-IOT/torch2trt
cd torch2trt
python setup.py install

这两个安装好之后就可以转模型了，首先的有一个.pth文件，就是ptorch训练的模型，我用的是mnist数据集训练测试的，训练如下

import torch
import numpy as np
from torch.nn import Module
from torch.autograd import Variable
import torch.nn as nn
from torch import nn,optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

batch_size = 64
learning_rate = 1e-2
num_epochs = 10

# 加载mnist数据集，构建训练集测试集的迭代器
data_tf = transforms.Compose([transforms.ToTensor(), transforms.Normalize([0.5], [0.5])])
train_dataset = datasets.MNIST(root='./data', train=True, transform=data_tf, download=True)
test_dataset = datasets.MNIST(root='./data', train=False, transform=data_tf)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

# cnn网络
class SimpleCnn(Module):
    def __init__(self):
        super(SimpleCnn, self).__init__() #b, 3, 32, 32
        self.layer1 = nn.Sequential()
        self.layer1.add_module('conv1', nn.Conv2d(1, 32,