Pytorch转Onnx转TensorRT

目前模型部署的常见做法是pytorch->onnx->tensorRT

首先使用pytorch框架训练深度学习模型，然后使用某些工具将训练好的pytorch模型转为onnx，最后转为tensorRT

获取pytorch模型

from torchvision import models

resnet18 = models.resnet18(pretrained=True)

将pytorch模型转成onnx模型

torch.onnx.export(model, torch.randn(1, 3, 224, 224), 'resnet18.onnx')

检查onnx模型

import onnx
import numpy as np
import onnxruntime as ort
import cv2

onnx_model = onnx.load('resnet18.onnx')
onnx.checker.check_model(onnx_model)

检查onnx模型是否能推理

inputs = np.random.randn(1, 3, 224, 224)
sess = ort.InferenceSession('resnet18.onnx', providers=ort.get_available_providers())
input_name = sess.get_inputs()[0].name
output = sess.run(None, {input_name:inputs})
print(output)

能够正常打印，即onnx模型没有问题

将onnx模型转成TensorRT模型

找到TensorRT-8.5.1.7文件（这是我自己的版本号）安装教程

将onnx文件移动到./TensorRT-8.5.1.7/bin 目录中

执行trtexec脚本，命令如下：

./trtexec --onnx=resnet18.onnx --saveEngine=resnet18.trt --workspace=6000

当前文件下会出现resnet18.trt 说明转换成功