Convert to ONNX and TensorRT

[pedrolsantos]

Hi,
Is it possible to convert this model to ONNX or to TensorRT ?
The architecture of this model seems a bit exotic and I'm having trouble converting it to ONNX...
Any experience or information regarding this type of conversion ?

[xiongxyowo]

Hi, the following demo script works fine in my environment, you can check if it works fine in yours as well:

import torch
from SAM2UNet import SAM2UNet
model = SAM2UNet().cuda()
model.eval()
input_names = ["input_0"]
output_names = ["output_0","output_1","output_2"]
x = torch.randn((1,3,352,352)).cuda()
torch.onnx.export(model,(x,),'sam2unet.onnx',input_names=input_names,output_names=output_names,
                    dynamic_axes={'input_0':[0],'output_0':[0]})
print("Done")

[pedrolsantos]

yes, this was very helpful
I managed to create the ONNX and exported to TensorRT
Got inference times around 11 ms with a 4060 TI
Thank you