pth to ONNX

[lhehejunl]

Hi,bigcow, can you provide the code which model.pth convert to model.onnx? thanks.

[MinGiSa]

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)