search

hmorimitsu / ptlflow

PyTorch Lightning Optical Flow models, scripts, and pretrained weights.
Apache License 2.0
250 stars 33 forks source link

export models to ONNX #26

Closed GotG closed 2 years ago

GotG commented 2 years ago

Hello

Do you have any suggestion/workflow on how to export the optical flow models to ONNX?

Thank you!

hmorimitsu commented 2 years ago

Hi,

Have you tried to follow the ONNX guide from PyTorch? I think it should work, although I am not sure if it can work for all models, as some of them may have some custom layers not supported by ONNX.

GotG commented 2 years ago

I loaded the models with ptlflow, then tried the conversion. unfortunately, none of the model conversions worked for me.

hmorimitsu commented 2 years ago

Hi,

I wrote a simple example to export to ONNX in this Colab: https://colab.research.google.com/drive/11xvIEerkxowM_lO-hEDNOqT6dANDZKUd?usp=sharing

Note, however, that it seems that ONNX does not support grid_sample at the moment, so many models cannot be exported, at least without changing their code.

GotG commented 2 years ago

Hello Thank you, looks good! I managed to convert flownets as well using ptl method. However the grid_sampler is a long running onnx support problem. Apparently it should be added to onnx (not pytorch yet) this November. Fingers crossed.