2、TensorRT学习笔记之PT转ONNX、可视化ONNX

NO郑

已于 2024-04-29 15:02:18 修改

阅读量1.4k

点赞数 10

分类专栏：模型部署学习笔记文章标签：计算机视觉笔记深度学习

于 2024-04-23 09:53:47 首次发布

本文链接：https://blog.csdn.net/NOZhengYuan/article/details/138110310

版权

模型部署学习笔记专栏收录该内容

4 篇文章 1 订阅

订阅专栏

本文详细介绍了如何从PyTorch模型（如YOLOv8）加载权重，关闭梯度更新，进行精度转换和融合操作，然后使用torch.onnx.export将模型转换为ONNX格式，并添加元数据。最后展示了完整的代码和ONNX可视化方法。

摘要由CSDN通过智能技术生成

摘要：主要是解析PT权重转ONNX的过程、代码。文章末尾附有完整运行代码。

目录：

1、导入模型

model  = torch.load("yolov8n.pt")  

# 由于我用v8做测试，v8的权重中除了模型结构，还有配置文件的参数，这里只需要其模型结构
ckpt = torch.load("yolov8n.pt")    
model = ckpt['model']

2、关闭梯度更新

作用：减少计算所需的资源

for p in model.parameters():
    p.requires_grad = False

3、精度转换、fuse等

fuse的作用：将conv和BN层融合，提高推理速度

model.eval()
model.float()
model = model.fuse()

4、参数的设置

device = 'cpu'        # 如果cuda能用，设置'cuda:0'
input = torch.zeros((1, 3, 640, 640)).to(device)      # 根据模型输入尺寸设置(1, 3, 640, 640)
f = 'name.onnx'                # 文件名，保存路径
opset_version = 11             # opset版本
input_names = ['img']          # 输入名
output_names=['out']           # 输出名
# 通过以下规则设置动态的维度
dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'},  # shape(1,3,640,640)
                         'output': {0: 'batch', 1: 'anchors'}  # shape(1,25200,85)}

5、 pt转onnx

torch.onnx.export(
    model，                            # 模型pt权重
    input,                             # 输入张量，模型输入，如[1,3,640,640]
    f,                                 # 保存onnx模型的文件
    opset_version = opset_version,     # Opset版本
    input_names=input_names,           # 输入张量名称
    output_names=output_names,         # 输出的张量名称
    dynamic_axes=dynamic_axes，        # 通过以下规则设置动态的维度
)

也可以结合步骤四直接这么写：

torch.onnx.export(model = torch.load(weight)[model],
                  im = torch.zeros(1, 3, 640, 640), 
                  f = ./weight.onnx, 
                  export_params = 12,
                  input_names=['images'],
                  output_names=['output'],
                  dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'}, 
                                'output': {0: 'batch', 1: 'anchors'}  # shape(1,25200,85)})

6、添加其余信息

比如你想在onnx中保存一些模型之外的信息，比如作者、模型名、classes、Opset版本等信息的话。

metadata = {
            'author': 'NO郑',
            'version': '11',
            'stride': int(max(model.stride)),
            'task': 'yolov8',
            'batch': 1,
            'imgsz': (640, 640),
            'labels': ['car', 'person']}  # model metadata
model_onnx = onnx.load(f)
for k, v in metadata.items():
    meta = model_onnx.metadata_props.add()
    meta.key, meta.value = k, str(v)

7、保存onnx

onnx.save(model, f)

8、完整代码

import torch
import onnx

ckpt = torch.load("yolov8n.pt")
model = ckpt['model']
im = torch.zeros((1, 3, 640, 640)).to('cpu')
for p in model.parameters():
    p.requires_grad = False

model.eval()
model.float()
model = model.fuse()
f = 'onnx01.onnx'

torch.onnx.export(model.cpu(), im, f, opset_version=11, input_names=['img'], output_names=['out'], dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'},  # shape(1,3,640,640)
                                'output': {0: 'batch', 1: 'anchors'}  # shape(1,25200,85)
                                })

metadata = {
            'description': '1',
            'author': 'NO郑',
            'version': '11',
            'batch': 1,
            'imgsz': (640, 640),
            'names': ['class1', 'classes2']}  # model metadata
model_onnx = onnx.load(f)
for k, v in metadata.items():
    meta = model_onnx.metadata_props.add()
    meta.key, meta.value = k, str(v)
onnx.save(model_onnx, f)