【DL】Pytorch_ONNX_ONNX-sim

最新推荐文章于 2024-09-14 16:52:58 发布

张林克

最新推荐文章于 2024-09-14 16:52:58 发布

阅读量1.5k

点赞数

分类专栏：【DeepL】

本文链接：https://blog.csdn.net/weixin_45292794/article/details/115653441

版权

本文详细介绍了如何将PyTorch模型转换为ONNX格式，并进行模型简化，适合嵌入式设备部署。内容包括PyTorch到ONNX的转换，ONNX模型的简化和可视化，以及生成可用于小站的模型和bin文件的代码示例。

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

把Pytorch模型转为ONNX，并进行精简化，用于嵌入式模型部署

Pytorch2ONNX

注意：多输出需要用 [out1, out2, out3] 列表进行存储

在这里插入图片描述

import os
import torch
import torchvision
from torchstat import stat
from torchsummary import summary
from yolo import YOLO


os.chdir('../')
ROOT = os.getcwd()

# 保存的文件名
ONNX_PATH = os.path.join(ROOT, "onnx", "xhh_yolov4_tiny.onnx")

# 权重文件
file_path = "logs/3.2 use_mosaic use_Cosine_lr use_anchors/Epoch600-Total_Loss1.6992-Val_Loss1.7980.pth"
MODEL_PATH = os.path.join(ROOT, file_path)

# [1]获取模型
def getYoloModel():
    from nets.yolo4_tiny import YoloBody

    # model = YoloBody(len(anchors[0]), len(class_names)).eval()
    model = YoloBody(3, 2, xhh_out_list_for_Inference=True).eval()

    print('Loading weights into state dict...')
    state_dict = torch.load(MODEL_PATH, map_location=