Open lhehejunl opened 1 year ago
from archs import UNext import torch
if name == 'main': with torch.no_grad(): device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = UNext(1,3,False).to(device) checkpointPath = "Your Model File Path" destPath = "Save Path, You Want" model.load_state_dict(torch.load(checkpointPath)) model.eval()
imgHeight = 256
imgWidth = 256
batchSize = 1
dummyInput = torch.rand(batchSize, 3, imgHeight, imgWidth).to(device)
inputNames = ["input"]
outputNames = ["output"]
dynamicAxes = {"input": {0: "batchSize"}, "output": {0: "batchSize"}}
torch.onnx.export(model,
dummyInput,
destPath,
input_names=inputNames,
output_names=outputNames,
dynamic_axes=dynamicAxes,
verbose=True)
Hi,bigcow, can you provide the code which model.pth convert to model.onnx? thanks.