yolov7-tiny.weigths转onnx模型

原创

已于 2023-06-01 16:39:46 修改 · 935 阅读

2 ·

CC 4.0 BY-SA版权

文章标签：

#YOLO #深度学习 #pytorch

于 2023-06-01 16:39:34 首次发布

该代码片段展示了如何使用Python和PyTorch库darknet2pytorch将YOLOv7-Tiny模型转换为ONNX模型。转换过程中，可以设置静态或动态批处理大小，并提供了导出ONNX模型的完整流程。

请注意：本仓库内的转换代码——仅支持yolov7-tiny不支持yolov7

需要借助的代码仓库

pytorch-YOLOv4/darknet2onnx.py at master · Tianxiaomo/pytorch-YOLOv4 · GitHubhttps://github.com/Tianxiaomo/pytorch-YOLOv4

import sys
import torch
from tool.darknet2pytorch import Darknet


def transform_to_onnx(cfgfile, weightfile, batch_size=1, onnx_file_name=None):
    model = Darknet(cfgfile)

    model.print_network()
    model.load_weights(weightfile)
    print('Loading weights from %s... Done!' % (weightfile))

    dynamic = False
    if batch_size <= 0:
        dynamic = True

    input_names = ["inp