How do you convert the pytorch model to onnx?

pango99 commented 3 years ago

Hi, How do you convert the pytorch model to onnx format? can you share the code?

thanks!

andylucny commented 3 years ago

Pytorch supports the ONNX format, thus it should be quite easy.

torch_model = modnet.module dummy = torch.randn(1, 3, 512, 672, requires_grad=True).to('cuda') torch.onnx.export(torch_model, dummy, "bgremover.onnx", export_params=True, opset_version=9, do_constant_folding=True)

Just in the case of the ModNet, it was not so easy. The conversion crashes on None values. They are a result of the fact that the ModNet has a different layout during training and inference. So I had to modify the ModNet implementation a bit. Instead of returning "None, None, values", we have to return just "values". And there is a problem also with the boolean parameter "inference" which we also have to remove from the forward() method header.

See _how_to_save_the_ModNet_model_asONNX/models/modnet.py, which I have added just now.

CardLin commented 3 years ago

I just use your code by click "s" to convert ModNet to ONNX. But it can't inference with OpenCV DNN module. Can you tell us which version of ModNet do you use? and which version of torch?

CardLin commented 3 years ago

I just upgrade my conda environment: conda install pytorch torchvision cudatoolkit=10.2 -c pytorch

Current version: pytorch==1.7.1 torchvision==0.8.2 torchsummary==1.5.1

And your convert code working great now!!

I think this issue can be closed!! Thanks @andylucny

andylucny / BackgroundRemovalByModNet4OpenCV

How do you convert the pytorch model to onnx? #2