pytorch model to onnx and then to tensorrt

prerequisites:

pytorch model file (xxx.pt)

export_onnx.py

Dependent library:

tensorrt 7.1.3

1. Convert torch model to onnx model

# 加载必要的依赖包
from modeling.deeplab import * 
import torch 
import torchvision 
import onnx 

# 定义加载原始模型的方法
def get_model(): 
    model = DeepLab(num_classes=4, backbone='mobilenet', output_stride=16,
    sync_bn=None, freeze_bn=False) 
    ckpt = torch.load('model_best.pth', map_location='cpu') 
    model.load_state_dict(ckpt['state_dict']) 
    model = model.cuda() 
    model.eval() 
    return model 

# 加载模型
model = get_model() 

# 配置onnx模型的节点名称，参数
input_names = ["input"] 
output_names = ["output"] 
example_tensor = torch.randn(1, 3, 320, 640, device='cuda') 
torch.onnx.export(model, example_tensor, "deeplab.onnx", verbose=True, 
opset_version=11, input_names=input_names, output_names=output_names) 

# 加载onnx模型 
onnx_model = onnx.load('deeplab.onnx') 

# 检查onnx模型是否正确导出
onnx.checker.check_model(onnx_model) 

2. Use the onnx model to convert the tensorrt model

1. After installing tensorrt, use tensorrt's tool trtexec to export the trt model.

#该命令不会修改模型的精度类型 
sudo ./trtexec --onnx=alexnet.onnx --saveEngine=alexnet.trt 
#该命令会修改模型的精度类型为int8
sudo ./trtexec --onnx=alexnet.onnx --saveEngine=alexnet.trt --int8 

3. Rewriting before and after processing

According to the image preprocessing method of the original algorithm and the postprocessing method of the model prediction results, the deployment and migration of the entire algorithm can be completed by rewriting it.

Fourth, pay attention to the details

1. Image preprocessing Try to use opencv to directly read the image and preprocess the image. (In general, it is necessary to rewrite the preprocessing method of torchvisison into a numpy method)

2. Post-processing of model prediction results. (The original algorithm is mainly an array of torch type, while the output result of the trt model is generally of numpy type, and the processing method needs to be changed from torch.tensor() to numpy type)

3. After the model is converted, the model prediction time is about 0.05s/frame, and the pre- and post-processing of the image is also time-consuming. Therefore, when the speed needs to be increased, the image pre-processing operation and the post-processing operation of the model prediction result need to be considered globally.