深度学习之格式转换笔记(一)：模型文件pt转onnx转tensorrt格式实操成功

原创

已于 2024-12-27 11:27:48 修改 · 1.1w 阅读

54 ·

CC 4.0 BY-SA版权

文章标签：

#深度学习 #python #pytorch

于 2021-01-22 12:35:02 首次发布

本文介绍了如何将PyTorch模型转换为ONNX格式，并详细说明了转换过程。接着，展示了如何使用TensorRT将ONNX模型加速，包括设置配置和读取引擎。重点在于模型转换与性能提升技术的应用。

在这里插入图片描述

文章目录

pt 转 onnx
onnx 转 tensorrt

pt 转 onnx

模型文件后缀说明

在深度学习中，常见的模型文件后缀有.pt、.pth 和.pkl。这些后缀名的模型文件本质上在格式上并无区别，只是不同开发者根据个人喜好选择使用不同的后缀。通常在保存模型文件时，常用的函数是 torch.save()。

保存与调用模型

保存

使用 torch.save(model.state_dict(), mymodel.pth) 语句来保存模型，这种方式只保存模型的权重参数，而不保存模型的结构。

调用

调用模型时，需要先重新构建模型结构，示例代码如下：

model = My_model(*args, **kwargs)  
pthfile = r'绝对路径' 
loaded_model = torch.load(pthfile, map_location='cpu')
model.load_state_dict(loaded_model['model'])
model.eval()

上述代码中，model.eval() 用于不启用 BatchNormalization 和 Dropout，确保在推理过程中不改变模型的权值。

pt 转 onnx 代码实现

from nn.mobilenetv3 import mobilenetv3_large,mobilenetv3_large_full,mobilenetv3_small
import torch
from nn.models import DarknetWithShh
from hyp import hyp

def convert_onnx():
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    model_path = 'weights/mbv3_large_75_light_final.pt' 
    backone = mobilenetv3_large(width_mult=0.75)
    model = DarknetWithShh(backone, hyp,light_head=True).to(device)

    model.load_state_dict(torch.load(model_path, map_location=device)['model'])

    model.to(device)
    model.eval()
    dummy_input = torch.randn(1, 3, 32, 32