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

首先需要安装两个必要的包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,