how to convert the pth file to onnx?

[youzipi]

how to convert the pth file to onnx?

[gaurav0651]

https://github.com/gaurav0651/dis-bg-remover/tree/main

[farazBhatti]

import torch
device='cpu'
input_shape = (1,3, 512, 512)

# Create a dummy input tensor with random values (you should use actual data)
sample_input = torch.randn(input_shape).to(device)
model_path = "isnet-general-full_model.pth"

# Load the traced model
model = torch.load(model_path,map_location=device)
# Ensure that the model is in evaluation mode
model.eval()
onnx_model_path='isnet-onnx-model.onnx'
torch.onnx.export(
    model,                  # PyTorch Model
    sample_input,                    # Input tensor
    onnx_model_path,        # Output file (eg. 'output_model.onnx')
    opset_version=12,       # Operator support version
    input_names=['input'] ,  # Input tensor name (arbitary)
    output_names=['output'] # Output tensor name (arbitary)