Export to ONNX: a better story

What are the better ways to convert models to ONNX?
I usually do:

  • trace model with torch.jit.trace()
  • use traced model with torch.onnx.export()
    This almost never works out the box with dynamic shapes.
    I have to decorate certain functions with @torch.jit.script.
    I have to remove certain libraries like einops and other stuff.
    I have to refactor a bunch of layers. For example, torch.stft, nn.TransformerEncoderLayer are not onnx-export able. The list of unsupported stuff is very long.
    How do we practically export stuff? Are there a set of guides or libraries we can use to help us?

Right now I am having issues because I need to export a model to ONNX and it has a bunch of einops operation. What approach has worked for you to rewrite this lines?