pytorch模型(.pth)转tensorrt模型(.engine)几种方式

序言

这篇文章总结了训练好的pytorch模型转成tensorrt模型部署的几种方式，转换原理流程大致如下：

1. 导出网络定义以及相关权重；
2. 解析网络定义以及相关权重；
3. 根据显卡算子构造出最优执行计划；
4. 将执行计划序列化存储；
5. 反序列化执行计划；
6. 进行推理

值得注意的是第三点，可以看到tensorrt转换出来的模型实际上是和硬件绑定的，也就是在部署的过程中，如果你的显卡和显卡相关驱动软件（cuda、cudnn）发生了改变，那么模型就得需要重新做转换。

一、trtexec

trtexec是在tensorrt包中自带的转换程序，该程序位于bin目录下，用起来比较方便，也是最简单的trt模型转换方式，在使用之前需要系统安装好cuda和cudnn，否则无法正常运行。使用示例如下：

首先将pytorch模型先转换成onnx模型，示例代码如下：

def torch2onnx(model_path,onnx_path):
    model = load_model(model_path)
    test_arr = torch.randn(1,3,32,448)
    input_names = ['input']
    output_names = ['output']
    tr_onnx.export(
        model,
        test_arr,
        onnx_path,
        verbose=False,
        opset_version=11,
        input_names=input_names,
        output_names=output_names,
        dynamic_axes={"input":{3:"width"}}            #动态推理W纬度，若需其他动态纬度可以自行修改，不需要动态推理的话可以注释这行
    )
    print('->>模型转换成功！')

trtexec转换命令如下：

固定尺寸模型转换：

./trtexec --onnx=repvgg_a1.onnx --saveEngine=repvgg_a1.engine --workspace=1024  --fp16

动态尺寸模型转换：

./trtexec --onnx=repvgg_a1.onnx --saveEngine=repvgg_a1.engine --workspace=1024 --minShapes=input:1x3x32x32 --optShapes=input:1x3x32x320 --maxShapes=input:1x3x32x640 --fp16

参数详解：

二、torch2trt

torch2trt是nvidia官方维护的一个易于使用的PyTorch到TensorRT转换器，使用起来也比较简单，但是相对于上面的方式环境配置更复杂一些，需要预先安装好torch、torch2trt、tensorrt，python环境下tensorrt安装方式示例，在Tensorrt的 .tar包中依次找到这几个whl包，直接使用pip安装即可：

#1、安装tensorrt
cd ~/TensorRT-8.2.4.2/python
pip install tensorrt-8.2.4.2-cp37-none-linux_x86_64.whl

#2、安装Python UFF wheel文件。只有当你将TensorRT与TensorFlow一起使用时才需要安装这个文件  用处：pb转tensorRT
cd ~/TensorRT-8.2.4.2/uff
pip install uff-0.6.9-py2.py3-none-any.whl

#3、安装Python graphsurgeon whl文件   用处：可以让TensorRT 自定义网络结构
cd ~/TensorRT-8.2.4.2/graphsurgeon
pip install graphsurgeon-0.4.5-py2.py3-none-any.whl

#注意trt7.0的版本没有这个包(不用装)
#4、安装Python onnx-graphsurgeon whl文件  
cd ~/TensorRT-8.2.4.2/onnx_graphsurgeon
pip install onnx_graphsurgeon-0.3.12-py2.py3-none-any.whl

#5、安装pycuda  可以通过它来实现python 下CUDA 的编程
pip install pycuda

#6、验证安装,打印出tensorrt版本，即安装成功
python
import tensorrt
tensorrt.__version__

torch2trt安装：

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

模型转换代码使用示例：

    model = load_model(model_path)
    model.cuda()
    arr = torch.ones(1, 3, 32, 448).cuda()
    model_trt = torch2trt(model,
                          [arr],
                          fp16_mode=True,
                          log_level=trt.Logger.INFO,
                          max_workspace_size=(1 << 32),
                          max_batch_size=1,
                          )
    torch.save(model_trt.state_dict(), os.path.join(output, "model_trt.pth"))

    logger.info("Converted TensorRT model done.")
    
    engine_file = os.path.join(output, "model_trt.engine")
    with open(engine_file, "wb") as f:
        f.write(model_trt.engine.serialize())
    logger.info("Converted TensorRT model engine file is saved for C++ inference.")

其中保存的pth模型可以用来加载到TRTModule中,TRTModule可以像torch模型一样正常推理：

from torch2trt import TRTModule

model_trt = TRTModule()

model_trt.load_state_dict(torch.load('model_trt.pth'))

而engine序列化文件可用于tensorrt程序加载推理，但是这里需要注意的是torch2trt不支持动态推理，更多使用示例参考torch2trt的github说明。

三、torch2trt dynamic

torch2trt dynamic是torch2trt的动态推理版，支持模型转换后的动态推理，使用起来和torch2trt基本一致，首先是安装：

git clone https://github.com/grimoire/torch2trt_dynamic.git 
cd torch2trt_dynamic
python setup.py develop

然后是使用示例，增加了一个动态尺度的参数：

from torch2trt_dynamic import torch2trt_dynamic
import torch
from torch import nn
from torchvision.models.resnet import resnet50
import os

# create some regular pytorch model...
model = resnet50().cuda().eval()

# create example data
x = torch.ones((1, 3, 224, 224)).cuda()

# convert to TensorRT feeding sample data as input
opt_shape_param = [
    [
        [1, 3, 128, 128],   # min
        [1, 3, 256, 256],   # opt
        [1, 3, 512, 512]    # max
    ]
]
model_trt = torch2trt_dynamic(model, [x], fp16_mode=False, opt_shape_param=opt_shape_param)
torch.save(model_trt.state_dict(), os.path.join(output, "model_trt.pth"))

 logger.info("Converted TensorRT model done.")
    
 engine_file = os.path.join(output, "model_trt.engine")
with open(engine_file, "wb") as f:
    f.write(model_trt.engine.serialize())
logger.info("Converted TensorRT model engine file is saved for C++ inference.")

四、parser解析onnx模型

如果不想借助工具转换，也可以自己编写代码，使用tensorrt的parser接口解析onnx模型，构建engine引擎，这种方式比较简单，不依赖其他库，并且支持动态推理模型转换，python代码示例如下：

# --*-- coding:utf-8 --*--
import pycuda.autoinit
import pycuda.driver as cuda
import tensorrt as trt
import time
import cv2, os
import numpy as np
import math

TRT_LOGGER = trt.Logger()

class HostDeviceMem(object):
    def __init__(self, host_mem, device_mem):
        """
        host_mem: cpu memory
        device_mem: gpu memory
        """
        self.host = host_mem
        self.device = device_mem

    def __str__(self):
        return "Host:\n" + str(self.host) + "\nDevice:\n" + str(self.device)

    def __repr__(self):
        return self.__str__()

def get_engine(max_batch_size=1, onnx_file_path="", engine_file_path="", fp16_mode=False, save_engine=False,input_dynamic=False):
    """
    params max_batch_size:      预先指定大小好分配显存
    params onnx_file_path:      onnx文件路径
    params engine_file_path:    待保存的序列化的引擎文件路径
    params fp16_mode:           是否采用FP16
    params save_engine:         是否保存引擎
    returns:                    ICudaEngine
    """
    # 如果已经存在序列化之后的引擎，则直接反序列化得到cudaEngine
    if os.path.exists(engine_file_path):
        print("Reading engine from file: {}".format(engine_file_path))
        with open(engine_file_path, 'rb') as f, \
                trt.Runtime(TRT_LOGGER) as runtime:
            return runtime.deserialize_cuda_engine(f.read())  # 反序列化
    else:  # 由onnx创建cudaEngine

        # 使用logger创建一个builder
        # builder创建一个计算图 INetworkDefinition
        explicit_batch = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
        # In TensorRT 7.0, the ONNX parser only supports full-dimensions mode, meaning that your network definition must be created with the explicitBatch flag set. For more information, see Working With Dynamic Shapes.

        with trt.Builder(TRT_LOGGER) as builder, \
                builder.create_network(explicit_batch) as network, \
                trt.OnnxParser(network, TRT_LOGGER) as parser:  # 使用onnx的解析器绑定计算图，后续将通过解析填充计算图
            # builder.max_workspace_size = 1 << 30  # 预先分配的工作空间大小,即ICudaEngine执行时GPU最大需要的空间
            config = builder.create_builder_config()
            config.max_workspace_size = 1 << 30

            builder.max_batch_size = max_batch_size  # 执行时最大可以使用的batchsize
            if fp16_mode:
                config.set_flag(trt.BuilderFlag.FP16)
            # builder.fp16_mode = fp16_mode

            # 解析onnx文件，填充计算图
            if not os.path.exists(onnx_file_path):
                quit("ONNX file {} not found!".format(onnx_file_path))
            print('loading onnx file from path {} ...'.format(onnx_file_path))
            with open(onnx_file_path, 'rb') as model:  # 二值化的网络结果和参数
                print("Begining onnx file parsing")
                parser.parse(model.read())  # 解析onnx文件
            # parser.parse_from_file(onnx_file_path) # parser还有一个从文件解析onnx的方法
            print("Completed parsing of onnx file")
            # 填充计算图完成后，则使用builder从计算图中创建CudaEngine
            print("Building an engine from file{}' this may take a while...".format(onnx_file_path))
            if input_dynamic:                              # 动态推理
                profile = builder.create_optimization_profile()
                profile.set_shape("input",(1,3,32,32),(1,3,32,320),(1,3,32,640))                
                config.add_optimization_profile(profile)
            #################
            print(network.get_layer(network.num_layers - 1).get_output(0).shape)
            engine =  builder.build_engine(network, config)
            print("Completed creating Engine")
            if save_engine:  # 保存engine供以后直接反序列化使用
                with open(engine_file_path, 'wb') as f:
                    f.write(engine.serialize())  # 序列化
            return engine

if __name__== "__main__":
    # These two modes are depend on hardwares
    fp16_mode = True
    max_batch_size = 1
    onnx_model_path = "./repvgg_a1.onnx"
    trt_engine_path = "./repvgg_a1.engine"
    # Build an cudaEngine
    engine = get_engine(max_batch_size, onnx_model_path, trt_engine_path, fp16_mode,True,True)

C++版本解析onnx代码示例：

//step1:创建logger：日志记录器
class Logger : public ILogger           
 {
     void log(Severity severity, const char* msg) override
     {
         // suppress info-level messages
         if (severity != Severity::kINFO)
             std::cout << msg << std::endl;
     }
 } gLogger;


//step2:创建builder
IBuilder* builder = createInferBuilder(gLogger);

//step3:创建network
const auto explicitBatch = 1U << static_cast<uint32_t>(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);  
INetworkDefinition* network = builder->createNetworkV2(explicitBatch);

//step4:创建parser
nvonnxparser::IParser* parser = nvonnxparser::createParser(*network, gLogger);

//step5:使用parser解析模型填充network
const char* onnx_filename="./model.onnx"
parser->parseFromFile(onnx_filename, ILogger::Severity::kWARNING);
for (int i = 0; i < parser.getNbErrors(); ++i)
{
    std::cout << parser->getError(i)->desc() << std::endl;
}

//step6:标记网络输出
for (auto &s : OUTPUT_BLOB_NAMES)
    network->markOutput(*blobNameToTensor->find(s.c_str()));

//step7:创建config并设置最大batchsize和最大工作空间
IBuilderConfig* config = builder->createBuilderConfig();
config->setMaxBatchSize(maxBatchSize);//设置最大batchsize
config->setMaxWorkspaceSize(1 << 30);//2^30 ,这里是1G

//step8:创建engine
ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
assert(engine);

//step9:序列化保存engine到planfile
IHostMemory *serializedModel = engine->serialize();
assert(serializedModel != nullptr)
std::ofstream p("xxxxx.engine");
p.write(reinterpret_cast<const char*>(serializedModel->data()), serializedModel->size());

//step10:释放资源
serializedModel->destroy();
engine->destroy();
parser->destroy()
network->destroy();
config->destroy();
builder->destroy();

五、tensorrtx

tensorrtx的模型构建方式比较奇特，首先使用tensorrt自带API将网络搭建起来，然后再将权重赋值进去，这种方式在转换的时候只要网络搭得没问题，基本上转换也不会问题，很好的解决了在onnx转trt过程中某些算子不支持的问题，但是过程比较复杂，况且不支持动态尺度推理，目前trt对onnx的支持很好，基本上onnx的模型都能转过去，所以这种方式如果不嫌麻烦的话可以试试，示例过程：

克隆tensorrtx

git clone https://github.com/wang-xinyu/tensorrtx.git

生成yolov5.wts文件，下载权重文件yolov5s.pt，将tensorrtx/yolov5/gen_wts.py复制到ultralytics/yolov5中，执行

python gen_wts.py

编译tensorrtx/yolov5,生成yolov5.engine文件

mkdir build
cd build
cmake ..
make

默认情况下，生成的是s模型和fp16推理。以及批次为1的engine，yolov5的其他模型可以在代码中修改相关参数，

#define USE_FP16
#define DEVICE 0                   // GPU ID
#define NMS_THRESH 0.4
#define CONF_THRESH 0.5
#define BATCH_SIZE 1

#define NET s            // s m x l 

copy文件yolov5.wts到tensorrtx/yolov5/build目录下，执行以下命令生成yolov5.engine

sudo ./yolov5 -s yolov5s.wts yolov5.engine s
sudo ./yolov5 -d yolov5s.engine ../samples

六、onnx-tensorrt

onnx-tensorrt是onnx官方的一个转换仓库，提供了诸多对应的Tensorrt版本的分支，比如我这里用的是8.2-EA，正确的编译方式为：

先将onnx-tensorrt git下来：

git clone --recursive -b 8.2-EA https://github.com/onnx/onnx-tensorrt.git

编译：

cd onnx-tensorrt
mkdir build
cd build
# /path/to/TensorRT-8.2.4.2改成自己的TensorRT绝对路径
cmake .. -DTENSORRT_ROOT=/path/to/TensorRT-8.2.4.2
make -j8
make install

编译完成后，如果你的cuda环境变量是配置好的，则不需要再配置，如果没有配置好则需要配置一下：

终端中输入：vim ~/.bashrc

#cuda
export PATH=/usr/local/cuda-11.4/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda-11.4/lib64:$LD_LIBRARY_PATH

保存退出，运行source ~/.bashrc刷新生效

onnx-tensorrt转换命令如下，序列化成engine引擎：

onnx2trt my_model.onnx -o my_engine.trt

转换为可读txt文本：

onnx2trt my_model.onnx -t my_model.onnx.txt

python端使用onnx-tensorrt：

#安装tensorrt
python3 -m pip install <tensorrt_install_dir>/python/tensorrt-8.x.x.x-cp<python_ver>-none-linux_x86_64.whl
#安装onnx
python3 -m pip install onnx==1.8.0
#安装onnx-tensorrt，在onnx-tensorrt目录下运行
python3 setup.py install

python代码使用示例，用来推理：

import onnx
import onnx_tensorrt.backend as backend
import numpy as np

model = onnx.load("/path/to/model.onnx")
engine = backend.prepare(model, device='CUDA:0')
input_data = np.random.random(size=(32, 3, 224, 224)).astype(np.float32)
output_data = engine.run(input_data)[0]
print(output_data)
print(output_data.shape)

七、Torch-TensorRT

Torch-TensorRT

八、手动解析ONNX（C++版）

Onnx2TensorRT

来源：三叔家的猫