Error jit trace Unet model#699

Unanswered

sivam25 asked this question inQ&A

sivam25

Dec 15, 2022

· 2 comments· 3 replies

Return to top

Discussion options

sivam25
Dec 15, 2022

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')

You must be logged in to vote

Replies: 2 comments 3 replies

Comment options

qubvel
Dec 15, 2022
Maintainer

Just call before tracingmodel.encoder.set_swish(memory_efficient=False)

You must be logged in to vote

1 reply

Comment options

sivam25 Dec 15, 2022
Author

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"

Comment options

butlerbrm
Jun 6, 2024

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?

You must be logged in to vote

2 replies

Comment options

qubvel Jun 6, 2024
Maintainer

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.