pytorch模型转TensorRT模型，并在jetson nano上加速推理

原创

已于 2023-07-08 20:21:41 修改 · 746 阅读

10 ·

CC 4.0 BY-SA版权

文章标签：

#python #深度学习 #pytorch

于 2023-07-08 20:18:03 首次发布

参考文章：PyTorch+TensorRT！20倍推理加速！
对其使用的脚本做了一定的修改。

pytorch2onnx

模型可视化使用了netron命令

安装命令：pip install netron

新建torch2onnx.py脚本

脚本和要转换的模型文件放在同一路径下，此案例使用MobileNetv2，可以改为其他模型，只需要正确导入模型即可

# coding:gbk
import torch
import torchvision
from torch.autograd import Variable
import netron
from model_v2 import MobileNetV2

print(torch.__version__)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# torch  -->  onnx
input_name = ['input']
output_name = ['output']
input = Variable(torch.randn(1, 3, 224, 224)).cuda()
# model = torch.load('zxq_MobileNetV2.pth', map_location="cuda:0")
model = MobileNetV2(num_classes=6).to(device)
    # load model weights
model_weight_path = "./zxq_MobileNetV2.pth"
model.load_state_dict(torch.load(model_weight_path, map_location=device))
torch.onnx.export(model, input, 'model_onnx.onnx',opset_version=12, input_names=input_name, output_names=output_name, verbose=True)
netron.start('model_onnx.onnx')

生成的onnx模型会出现在同一路径下。

onnx2trt

同一路径下新建脚本onnx2trt.py

import tensorrt as trt


def build_engine(onnx_file_path,engine_file_path,half=False):
    """Takes an ONNX file and creates a TensorRT engine to run inference with"""
    logger = trt.Logger(trt.Logger.INFO)
    builder = trt.Builder(logger)
    config = builder.create_builder_config()
    config.max_workspace_size = 4 * 1 <