How to best save optimzed_for_inference models ?

[MateoPetitet]

Hello !

I would like to optimize my models for inference using the built-in method. However, I don't understand how to save it in a good way. All examples I saw optimize the model on-the-fly, before starting inference, which seems very unoptimized. I would like to have an optimized version to directly share it and use as-is, is it possible ?

Thank you all for your help !

[isaacrob-roboflow]

That method should only take a few seconds so we didn't think anyone would care to cache that work 😅 but it's just a JIT scripted model, there are native PyTorch tools for saving / loading them. Feel free to submit a PR for it if you make a nice pipeline

[MateoPetitet]

Oh really it is supposed to be fast ? I tried to convert before using it because I wanna use my model on 4GB of VRAM. I use a fine-tunned Large variant, which fits on ~2.3GB, but the conversion seems to use a lot of VRAM and I always end up with a CUDA low-memory error ! Hence why I wanna save directly.

[isaacrob-roboflow]

Oh interesting, I guess I didn't check the memory usage of the method. Are you able to run a predict at all with that much VRAM? I am not sure the optimize method will meaningfully decrease VRAM requirements

[MateoPetitet]

Are you able to run a predict at all with that much VRAM?

Yep, it work and I manage to do inference at a speed of ~1.5s/image. I don't have a requirement of real time for my task so it's good, but if I can speed it up using the optimization (and maybe do quantization to 16 bit, as I run it on a T600) I'll take it !

[isaacrob-roboflow]

Ok. Well in that case it seems like saving would be useful for you. It's a torchscript object, should be lots of documentation online about how to do it. For fp16, try setting dtype to float16 in the optimize call. Let me know if that works for you, may be precision issues because of a different bug