search

Z-Xiong / OSTrack-mnn

Mnn version demo of [ECCV2022] Joint Feature Learning and Relation Modeling for Tracking: A One-Stream Framework.
16 stars 2 forks source link

How to export ONNX model from pytorch code? #4

Closed Z-Xiong closed 2 months ago

Z-Xiong commented 3 months ago 
          Hello, could you share your PyTorch to onnx conversion code piece? As I understand ostrack-256-ep300.onnx does not contain LayerNorm() which I intend to do with my model as well (due to an old tensorrt version). It would be really useful. I appreciate any help you can provide.

Originally posted by @ozgyal in https://github.com/Z-Xiong/OSTrack-mnn/issues/2#issuecomment-2307164993

Z-Xiong commented 3 months ago

lib.zip torch2onnx.zip

ozgyal commented 3 months ago

Hello again, thank you so much for your quick response and sharing your code that fast. This helped a lot.

I've got only one question about precision casting. The ostrack-256-ep300.onnx model that you have shared in google drive link before includes a module named "Cast" between Sub and Pow modules (I've checked the model structure with netron.app)

I couldn't export this model with the code you have provided. Did you do some post-processing over the onnx model, probably for FP32 casting? If you remember something like that it would be great to have that post-process step as well.

Thank you in advance.

Z-Xiong commented 3 months ago

Thank you for your feedback. Last time I uploaded the code for converting to onnx, but the one in Google Drive is the onnx model needed for converting to mnn, they are not the same.

Dowdload pretrained-model to the project root path: output/checkpoints/train/ostrack/vitb_256_mae_ce_32x4_ep300/OSTrack_ep0300.pth.tar

torch2any.zip Try running the code torch2mnn.py as follows: python torch2mnn.py ostrack vitb_256_mae_ce_32x4_ep300 You will get the onnx model same as the model in google drive.

If you have any further questions, feel free to reach out again.