Unable to export Pytorch Deeplab model to ONNX for ORTModule Usage

[zazhil-ha]

Hello,

I am trying to use Pytorch Deeplab 3 model with ORTModule. However, when the code tries to export the model to ONNX format I get the following error:

ValueError: Invalid item in shape: floor(ATen_13_o0__d2/2 - 1/2) + 1. Needs to be of int or str.

Debugging I found out that sympy.core.add.Add is unknown to onnx so I tried to register it as a custom operation with the following code:

register_custom_op_symbolic('::add', symbolic_opset14.simpy_add, 14)

The definition in opset14 goes like:

@_beartype.beartype
def simpy_add(g, self, *args):
output=g.op("sympy.core.add::Add", self, *args)
output.setType(self.type())
return output

By doing this a new problem arises:

File ~/anaconda3/envs/py310ort/lib/python3.10/site-packages/onnxruntime/tools/symbolic_shape_infer.py:32, in is_sequence(type_proto)
30 def is_sequence(type_proto):
31 cls_type = type_proto.WhichOneof("value")
---> 32 assert cls_type in ["tensor_type", "sequence_type"]
33 return cls_type == "sequence_type"

AssertionError:

Seems like the result of the Add operation has no type. Could you please help me to properly define sympy function to make this work.

I am currently using Pyhon 3.10, onnx 1.14.1 and Pytorch 2.0.1+cu117

Thanks,

Zazhil-ha

[skottmckay]

Based on the documentation sympy.core.add.Add is not returning a value if you look at the Examples here: https://docs.sympy.org/dev/modules/core.html#module-sympy.core.add

I'm not familiar with sympy but I assume there's some other step to actually execute the expression. I would expect that's what your custom op implementation probably needs to do so it's returning a value that can be expressed as a tensor or sequence of tensors rather than returning the sympy Add class instance.

Also note that exporting is one half of the story. You would need an implementation of the custom op in C++ for ONNX Runtime to be able to load and execute the model.

https://github.com/onnx/tutorials/blob/main/PyTorchCustomOperator/README.md may be helpful in general.

[zazhil-ha]

I will give it a try, thanks!