1、pth转onnx模型、onnx转tensorrt模型、python中使用tensorrt进行加速推理（全网最全，不信你打我）

本文向所有亲们介绍在python当中配置tensorrt环境、使用tensorrt环境进行推理的教程，主要分为两大部分，第一部分环境配置，第二部分前向推理。

第一部分 环境配置

第一步：检查你的系统类型、cuda、cudnn版本，其中cuda查看命令及本人环境显示如下：

    nvcc -V

显示如下

nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2021 NVIDIA Corporation
Built on Wed_Jun__2_19:15:15_PDT_2021
Cuda compilation tools, release 11.4, V11.4.48
Build cuda_11.4.r11.4/compiler.30033411_0

查看ubuntu系统信息命令：

sudo lsb_release -a

显示如下：

o LSB modules are available.
Distributor ID:	Ubuntu
Description:	Ubuntu 18.04.6 LTS
Release:	18.04
Codename:	bionic

第二步： 官方下载链接下载tensorrt包。注意，需要亲亲注册哦，不然无法下载的啦。本人下载的是tar格式的

第三步：解压文件：

tar -xvf TensorRT-8.4.3.1.Linux.x86_64-gnu.cuda-11.6.cudnn8.4.tar.gz

第四部：激活你的当前python环境的情况下，进入你的解压包，安装对应python版本的tensorrt文件。在python文件夹下的内容如下，可以看到有好几个python对应版本的tensorrt，选择亲亲当前环境下的python版本进行安装就可以了，我的python版本是3.8，我就安装tensorrt-8.4.3.1-cp38-none-linux_x86_64.whl。

cd TensorRT-8.4.3.1/python/
pip install tensorrt-8.4.3.1-cp38-none-linux_x86_64.whl

第五步： 这下亲亲是不是觉得已经安装完成了可以使用了，当然不是，此时，一旦你在python中导入tensorrt时，会显示找不到包了的。所以，亲亲们要把解压出来的tensorrt文件当中的lib添加到环境变量当中哦！

export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/opt/TensorRT-8.4.3.1/lib

当然了，你还可以把lib和include中的所有文件复制到系统路径下面，命令如下：

sudo cp -r /opt/TensorRT-8.4.3.1/lib/* /usr/lib
sudo cp -r /opt/TensorRT-8.4.3.1//include/* /usr/include

第六步：检查python中tensorrt是否安装成功

>>> import tensorrt as tr
>>> tr.__version__
'8.4.3.1'
>>> 

至此，python中配置tensorrt环境就完成了（码字不易，赶紧点赞，别想白嫖）

第二部分 模型转化

第一步：将pytorch中训练出来的pth模型转化为onnx,转化脚本如下：

import torch
import onnx
def pth_to_onnx(input_path,output_path):
    '''
    1)声明：使用本函数之前，必须保证你手上已经有了.pth模型文件.
    2)功能：本函数功能四将pytorch训练得到的.pth文件转化为onnx文件。
    '''
    torch_model = torch.load(input_path)       # pytorch模型加载,此处加载的模型包含图和参数
    #torch_model = selfmodel()  # 若只需要保存参数，可以换成这一种，其中selfmodel需要自己编写
    torch_model.eval()
    x = torch.randn(1,1,28,28)          # 输入一张28*28的灰度图像并生成张量
    export_onnx_file = output_path         #输出.onnx文件的文件路径及文件名
    torch.onnx.export(torch_model,
                      x,
                      export_onnx_file,
                      opset_version=9,    #操作的版本，稳定操作集为9
                      do_constant_folding=True,          # 是否执行常量折叠优化
                      input_names=["input"],        # 输入名
                      output_names=["output"],       # 输出名
                      dynamic_axes={"input": {0: "batch_size"},         # 批处理变量
                                    "output": {0: "batch_size"}}
                      )

    #onnx_model = onnx.load('model_all.onnx')    #加载.onnx文件
    #onnx.checker.check_model(onnx_model)
    #print(onnx.helper.printable_graph(onnx_model.graph))       #打印.onnx文件信息

第二步：将onnx格式的转化为tensorrt支持的格式，转化方法比较多，这里只写一种在终端当中的方法（说实话，其他方法我也没有试哈哈）：

（1）进入到之前安装的Tensorrt安装路径，该路径下有一个trtexec文件

cd /opt/TensorRT-8.4.3.1/bin

（2）使用最简单的命令转化格式：

sudo ./trtexec --onnx=/home/photo-de/mm.onnx --saveEngine=m1.engine --fp16

其中，–onnx后面是你的onnx文件的路径， --saveEngine后面跟的是生成的engine文件的保存路径。–fp16 指生成半精度 模型，加快推理速度。如果只需要简单的转化以下就做到这一步可以了，否则请看（3）
（3）./trtexec 后面有很多参数，可以根据自己的实际情况来设置相应的参数。如下（这部分是引用）：

#静态输入尺度
./trtexec     --onnx=<onnx_file> \                      
            --explicitBatch \                             #在构建引擎时使用显式批大小(默认=隐式)显示批处理
            --saveEngine=<tensorRT_engine_file> \         #输出engine
            --workspace=<size_in_megabytes> \             #设置工作空间大小单位是MB(默认为16MB)
            --fp16                                         #除了fp32之外，还启用fp16精度(默认=禁用)
        
#动态输入尺度
./trtexec     --onnx=<onnx_file> \                       
            --minShapes=input:<shape_of_min_batch> \     #最小的batchsize x 通道数 x 输入尺寸x x 输入尺寸y
            --optShapes=input:<shape_of_opt_batch> \      #最佳输入维度，跟maxShapes一样就好
            --maxShapes=input:<shape_of_max_batch> \     #最大输入维度
            --workspace=<size_in_megabytes> \             #设置工作空间大小单位是MB(默认为16MB)
            --saveEngine=<engine_file> \                   #输出engine
            --fp16                                       #除了fp32之外，还启用fp16精度(默认=禁用)

举个例子：

#小尺寸的图片可以多batchsize即8x3x416x416
/home/zxl/TensorRT-7.2.3.4/bin/trtexec  --onnx=mymodel.onnx \
                                        --minShapes=input:1x3x416x416 \
                                        --optShapes=input:8x3x416x416 \
                                        --maxShapes=input:8x3x416x416 \
                                        --workspace=4096 \
                                        --saveEngine=mymodel.engine \
                                        --fp16
 
#由于内存不够了所以改成4x3x608x608
/home/zxl/TensorRT-7.2.3.4/bin/trtexec  --onnx=mymodel.onnx \
                                        --minShapes=input:1x3x608x608 \
                                        --optShapes=input:4x3x608x608 \
                                        --maxShapes=input:4x3x608x608 \
                                        --workspace=4096 \
                                        --saveEngine=mymodel.engine \
                                        --fp16  

到此，我们得到了最终的engine文件。

第三部分 使用python+tensorrt进行模型推理

在你的推理环境当中需要再次安装pycuda，numpy。直接上代码：

class TrtModel:
    def __init__(self, engine_path, max_batch_size=1, dtype=np.float32):

        self.engine_path = engine_path
        self.dtype = dtype
        self.logger = trt.Logger(trt.Logger.WARNING)
        self.runtime = trt.Runtime(self.logger)
        self.engine = self.load_engine(self.runtime, self.engine_path)
        self.max_batch_size = max_batch_size
        self.inputs, self.outputs, self.bindings, self.stream = self.allocate_buffers()
        self.context = self.engine.create_execution_context()

    @staticmethod
    def load_engine(trt_runtime, engine_path):
        trt.init_libnvinfer_plugins(None, "")
        with open(engine_path, 'rb') as f:
            engine_data = f.read()
        engine = trt_runtime.deserialize_cuda_engine(engine_data)
        return engine

    def allocate_buffers(self):

        inputs = []
        outputs = []
        bindings = []
        stream = cuda.Stream()

        for binding in self.engine:
            size = trt.volume(self.engine.get_binding_shape(binding)) * self.max_batch_size
            host_mem = cuda.pagelocked_empty(size, self.dtype)
            device_mem = cuda.mem_alloc(host_mem.nbytes)

            bindings.append(int(device_mem))

            if self.engine.binding_is_input(binding):
                inputs.append(HostDeviceMem(host_mem, device_mem))
            else:
                outputs.append(HostDeviceMem(host_mem, device_mem))
        return inputs, outputs, bindings, stream

    def __call__(self, x: np.ndarray, batch_size=2):

        x = x.astype(self.dtype)

        np.copyto(self.inputs[0].host, x.ravel())

        for inp in self.inputs:
            cuda.memcpy_htod_async(inp.device, inp.host, self.stream)

        self.context.execute_async(batch_size=batch_size, bindings=self.bindings, stream_handle=self.stream.handle)
        for out in self.outputs:
            cuda.memcpy_dtoh_async(out.host, out.device, self.stream)

        self.stream.synchronize()
        return [out.host.reshape(batch_size, -1) for out in self.outputs]
def predicts(model_path,images):
    batch_size = 1
    model = TrtModel(model_path)
    result = model(images, batch_size)
    return result
if __name__ == '__main__':
    model_path = 'you model '
    images = 'you image'
    pred = predicts(model_path,images)