Error jit trace Unet model

[sivam25]

I am trying to jit trace the unet model. But I get an error saying "Could not export Python function call 'SwishImplementation'. Remove calls to python functions before export." How can I export the model so I can use it in a c++ application.

model = smp.Unet('efficientnet-b7')
model.eval()

input = torch.randn((1,3,224,224))
torch_out = model(input)

model = torch.jit.trace(model,input)
trace_out = model(input)

model.save('model.pt')

[qubvel]

Just call before tracing model.encoder.set_swish(memory_efficient=False)

[sivam25]

Hi, Thanks that fixed the issue. But im unable to load the model.pt file in c++ using torch::jit::load(). Also how can I export the model to a onnx format? I Keep getting warning saying "Constant folding- only steps=1 can be folded for opset >=10"

[butlerbrm]

I had an issue with this case as well. tho after using Just call before tracing model.encoder.set_swish(memory_efficient=False). I got this
segmentation_models_pytorch\base\model.py:16: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
if h % output_stride != 0 or w % output_stride != 0:

Any suggestions?

[qubvel]

The model is checking that the input size is divisible by 32 or 16 to give a proper error. In the case of tracing it should be fine, but shape-checking will not work in the traced model.

[butlerbrm]

I see. Thank you