u2net导onnx

最新推荐文章于 2024-07-17 17:50:29 发布

Alphapeople

最新推荐文章于 2024-07-17 17:50:29 发布

阅读量441

点赞数 3

文章标签：深度学习 pytorch 人工智能

本文链接：https://blog.csdn.net/weixin_38241876/article/details/137682535

版权

import torch
import torchvision as tv
from u2net import U2NET

def deeplabv32onnx():
    model = U2NET()
    model.load_state_dict(torch.load(r'D:\Dely\unet.pt'))
    model.eval()
    dummy_input = torch.randn(1, 3, 512,512)
    model(dummy_input)
    im = torch.zeros(1, 3,512,512)#.to("cpu")
    torch.onnx.export(model, im,
                      "u2net.onnx",
                      verbose=False,
                      opset_version=11,
                      training=torch.onnx.TrainingMode.EVAL,
                      do_constant_folding=True,
                      input_names=['input1'],
                      output_names=['o0', 'o1','o2','o3','o4','o5','o6'],
                      dynamic_axes={'input': {0: 'batch', 2: 'height', 3: 'width'}}
                      )
deeplabv32onnx()