为了方便在C#项目中引用onnx文件,于是需要将pth模型文件转换为onnx类型。
转换的模型项目地址是:https://github.com/xuebinqin/U-2-Net,以下为python的示例代码:
1 import torch
2 import sys 3 import os 4 model_dir=os.path.join(os.path.dirname(__file__), 'model') 5 sys.path.append(model_dir) 6 from model import U2NET 7 net = U2NET(3,1) 8 net.load_state_dict(torch.load("u2net.pth", map_location='cpu')) 9 net.eval() 10 # 设置输入样例 11 dummy_input = torch.randn(1, 3, 320, 320) # 根据实际输入大小调整 12 # 导出模型到ONNX格式 13 torch.onnx.export(net, dummy_input, "u2net.onnx", verbose=True, input_names=['input_1'], output_names=['output_1'])
以下是当前文件夹下相关文件信息:
标签:转换,pth,python,onnx,torch,input,import,net From: https://www.cnblogs.com/yu-xiansheng/p/18289796