I want to export a PyTorch model to ONNX using torch.onnx.export and I have some custom operators in it. I have managed to add them to TorchScript’s operator registry and I export the model fine to ONNX and Netron shows everything is ok.
The only issue is I am getting this warning: WARNING: The shape inference of custom_library::custom_op type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function.
Now I’d love to add some shape inference but I do not now how to. From the relevant docs it seems it needs to be done using the .setType function on the call to g.op in the symbolic function definition. Despite that, .setType receives only types and not dims/shapes.
Furthermore, I couldn’t find any mention of shape inference in Python not in PyTorch’s docs nor in ONNX’s docs. What is the best way to do this then?
I’ve experienced the same issue, let me share what I’ve found. Maybe it will be helpful.
Here’s the #73871 pytorch issue for documenting shape inference for custom symbolics:
Here’s the example that they are referring to.
But at the same time there is also #81693 pytorch issue saying that even with invoking g.op(...).setType() inside the symbolic function it is not sufficient to make shape inference work.
Seems like it’s going to be fixed in next release of torch: PR
I’ve tried latest nightly build torch==1.14.0.dev20221119+cu116, but it’s still not working.
The commit was merged yesterday, possibly it should be in the next build.