Varias formas de convertir el modelo pytorch (.pth) al modelo tensorrt (.engine)

Prefacio

Este artículo resume varias formas de convertir un modelo de Pytorch entrenado en un modelo de tensorrt para su implementación. El proceso de principio de conversión es aproximadamente el siguiente:

1. Exportar definiciones de red y pesos relacionados;
2. Analizar definiciones de red y pesos relacionados;
3. Construir el plan de ejecución óptimo según el operador de la tarjeta gráfica;
4. Serializar y almacenar el plan de ejecución;
5. Deserializar el plan de ejecución;
6. hacer inferencias

Vale la pena señalar el tercer punto: puede ver que el modelo convertido por tensorrt en realidad está vinculado al hardware, es decir, durante el proceso de implementación, si su tarjeta gráfica y el software del controlador relacionado con la tarjeta gráfica (cuda, cudnn) cambian, entonces será necesario volver a convertir el modelo.

1. trtexec

trtexec es un programa de conversión que viene con el paquete tensorrt. Este programa se encuentra en el directorio bin. Es fácil de usar y es la forma más sencilla de convertir modelos trt. Antes de usarlo, es necesario instalar cuda y cudnn en el sistema. de lo contrario no funcionará normalmente. Los ejemplos de uso son los siguientes:

Primero, convierta el modelo pytorch en un modelo onnx, el código de muestra es el siguiente:

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('->>模型转换成功！')

El comando de conversión trtexec es el siguiente:

Conversión de modelo de tamaño fijo:

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

Conversión de modelo de tamaño dinámico:

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

Explicación detallada de los parámetros:

• –onnx ruta onnx
• –saveEngine trt dirección de guardado del motor de inferencia de serialización
• –workspace Establece el tamaño del espacio de trabajo en megabytes (predeterminado = 16)
• –minShapes Genera formas dinámicas utilizando el archivo de configuración de forma mínima proporcionado
• –optShapes Genera formas dinámicas utilizando el archivo de configuración de formas óptimas proporcionado
• –maxShapes Genera formas dinámicas utilizando el archivo de configuración de forma máxima proporcionado
• –fp16 Activa la inferencia de precisión float16 (se recomienda este modo, por un lado puede acelerar, por otro lado la disminución de la precisión es relativamente pequeña)

2. antorcha2trt

torch2trt es un conversor de PyTorch a TensorRT fácil de usar mantenido oficialmente por nvidia. Es relativamente simple de usar, pero la configuración del entorno es más complicada que el método anterior. Torch, torch2trt y tensorrt deben instalarse con anticipación. para instalar tensorrt en el entorno python. Por ejemplo, busque estos paquetes whl en el paquete .tar de Tensorrt e instálelos directamente usando 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__

Instalación de torch2trt:

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

Ejemplo de uso del código de conversión de modelo:

    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.")

El modelo pth guardado se puede cargar en TRTModule, y TRTModule puede razonar normalmente como el modelo de antorcha:

from torch2trt import TRTModule

model_trt = TRTModule()

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

El archivo serializado del motor se puede usar para la inferencia de carga del programa tensorrt, pero debe tenerse en cuenta aquí que torch2trt no admite la inferencia dinámica. Para obtener más ejemplos de uso, consulte las instrucciones de github de torch2trt .

三、Torch2trt dinámico

torch2trt dinámico es la versión de razonamiento dinámico de torch2trt, que admite el razonamiento dinámico después de la conversión del modelo. En uso, es básicamente lo mismo que torch2trt. El primer paso es instalar:

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

Luego hay un ejemplo de uso, agregando un parámetro de escala dinámica:

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.")

4. El analizador analiza el modelo onnx

Si no desea utilizar la conversión de herramientas, también puede escribir su propio código, usar la interfaz del analizador de tensorrt para analizar el modelo onnx y construir el motor. Este método es relativamente simple, no depende de otras bibliotecas y admite Conversión del modelo de inferencia dinámica. El ejemplo del código Python es el siguiente:

# --*-- 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)

Ejemplo de código onnx de análisis de la versión C++:

//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();

5. tensorrx

El método de construcción del modelo de tensorrtx es bastante extraño. Primero, use la propia API de tensorrt para construir la red y luego asigne los pesos. De esta manera, siempre que la red esté bien establecida durante la conversión, básicamente la conversión no será un problema. , lo cual es bueno. Resuelve el problema de que algunos operadores no son compatibles en el proceso de conversión de onnx a trt, pero el proceso es relativamente complicado y el razonamiento de escala dinámica no es compatible. Actualmente, trt admite muy bien onnx y básicamente todo Los modelos onnx se pueden convertir, así que si este método no es demasiado problemático, puedes probarlo. Proceso de ejemplo:

Clonar tensorrtx

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

Genere el archivo yolov5.wts, descargue el archivo de peso yolov5s.pt, copie tensorrtx/yolov5/gen_wts.py a ultralytics/yolov5 y ejecute

python gen_wts.py

Compile tensorrtx/yolov5 y genere el archivo yolov5.engine

mkdir build
cd build
cmake ..
make

De forma predeterminada, se generan la inferencia s-model y fp16. Además del motor del lote 1, otros modelos de yolov5 pueden modificar parámetros relevantes en el código.

#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 

Copie el archivo yolov5.wts al directorio tensorrtx/yolov5/build y ejecute el siguiente comando para generar yolov5.engine

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

6. onnx-tensorrt

onnx-tensorrt es un almacén de conversión oficial de onnx, que proporciona muchas ramas correspondientes de la versión Tensorrt. Por ejemplo, aquí uso 8.2-EA. El método de compilación correcto es:

Primero baja onnx-tensorrt:

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

Compilar:

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

Una vez completada la compilación, si sus variables de entorno cuda están configuradas, no necesita configurarlas nuevamente. Si no están configuradas, debe configurarlas:

Ingrese en la terminal: 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

Guarde y salga, ejecute source ~/.bashrc para actualizar y surtir efecto.

El comando de conversión onnx-tensorrt es el siguiente, serializado en el motor:

onnx2trt my_model.onnx -o my_engine.trt

Convertir a texto de texto legible:

onnx2trt my_model.onnx -t my_model.onnx.txt

Utilice onnx-tensorrt en el lado de Python:

#安装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

Ejemplo de uso del código Python para razonar:

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

Antorcha-TensorRT

8. Analizar manualmente ONNX (versión C++)

Onnx2TensorRT