ONNX export fails for TFSegformerForSemanticSegmentation

[OutSorcerer]

System Info

• transformers version: 4.27.0.dev0
• Platform: Linux-5.10.147+-x86_64-with-glibc2.29
• Python version: 3.8.10
• Huggingface_hub version: 0.12.0
• PyTorch version (GPU?): 1.13.1+cu116 (False)
• Tensorflow version (GPU?): 2.11.0 (False)
• Flax version (CPU?/GPU?/TPU?): not installed (NA)
• Jax version: not installed
• JaxLib version: not installed
• Using GPU in script?: No
• Using distributed or parallel set-up in script?: No

Additionally, the versions of some relevant packages are

transformers @ git+https://github.com/huggingface/transformers@762dda44deed29baab049aac5324b49f134e7536
onnx==1.13.0
onnxruntime==1.14.0
tf2onnx==1.13.0

Who can help?

@gante, @Rocketknight1

Information

• [ ] The official example scripts
• [ ] My own modified scripts

Tasks

• [ ] An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• [ ] My own task or dataset (give details below)

Reproduction

1. Run this notebook (https://github.com/deep-diver/segformer-tf-transformers/blob/main/notebooks/TFSegFormer_ONNX.ipynb) in Colab. ONNX export apparently worked there as of July 25 2022, but it fails now.

   The error message is

ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.0.0/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.0.1/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.0.2/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.1.0/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.1.1/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.1.2/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.1.3/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.1.4/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.1.5/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.2.0/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.2.1/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.2.2/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.2.3/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.2.4/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.2.5/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.2.6/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.2.7/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.2.8/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.2.9/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.2.10/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.2.11/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.2.12/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.2.13/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.2.14/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.2.15/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.2.16/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.2.17/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.2.18/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.2.19/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.2.20/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.2.21/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.2.22/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.2.23/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.2.24/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.2.25/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.2.26/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.2.27/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.2.28/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.2.29/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.2.30/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.2.31/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.2.32/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.2.33/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.2.34/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.2.35/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.2.36/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.2.37/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.2.38/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.2.39/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.3.0/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.3.1/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [tf_segformer_for_semantic_segmentation/segformer/encoder/block.3.2/mlp/dwconv/dwconv/PartitionedCall: PartitionedCall] is not supported
ERROR:tf2onnx.tfonnx:Unsupported ops: Counter({'PartitionedCall': 52})

Expected behavior

ONNX export of TFSegformerForSemanticSegmentation works.

[Rocketknight1]

cc @sayakpaul, any idea what the cause might be?

[sayakpaul]

Maybe this issue is better suited for https://github.com/onnx/tensorflow-onnx

Could you try downgrading the tf2onnx version to 1.11.1?

[OutSorcerer]

@sayakpaul

Could you try downgrading the tf2onnx version to 1.11.1?

I tried this, unfortunately there is still an error that PartitionedCall is not supported.

I also tried to downgrade tensorflow and ONNX export worked with tensorflow==2.8.4, here is an example: https://colab.research.google.com/gist/OutSorcerer/c8cd27a455091b57d9ea90ab3450035e/tfsegformer_onnx.ipynb

Maybe this issue is better suited for https://github.com/onnx/tensorflow-onnx

There are already issues there about PartitionedCall support e.g. https://github.com/onnx/tensorflow-onnx/issues/1864.

However, since export works with a previous version of TensorFlow, it seems that PartitionedCall operation is not essential for a model to work. This is a low-level operation automatically added by TensorFlow and another workaround with new versions of TensorFlow could be to disable its insertion into an operation graph, but I was not able to quickly find a way to do it.

Also, regardless of the error message printed an ONNX file is still generated (which obviously fails at inference time), so yet another workaround could be to remove PartitionedCalls from an ONNX file.

[sayakpaul]

Thanks for investigating. With your workaround, does the model work during inference as expected?

If so, I guess we can safely close the issue here?

[OutSorcerer]

Thanks for investigating. With your workaround, does the model work during inference as expected?

Yes, I rerun the cells that were comparing outputs of a TF model and an ONNX model and the outputs match.

If so, I guess we can safely close the issue here?

Well, from my perspective ideally one of workarounds would be applied in transformers and TFSegformerForSemanticSegmentation would work with the most recent releases of TF and other packages, but I also understand that eventually tf2onnx developers should do something with PartitionedCall export and this issue would be solved too.

[OutSorcerer]

In fact, PartitionedCall may not be the root cause of the problem.

I looked at the ONNX file produced with TF 2.11.0 in the notebook above by doing

onnx_model = onnx.load(onnx_model_path)
with open("model.txt", "w") as f:
  f.write(str(onnx_model))

It has the following node

node {
    input: "tf_segformer_for_semantic_segmentation/segformer/encoder/block.1.0/mlp/dwconv/Reshape:0"
    input: "tf_segformer_for_semantic_segmentation/segformer/encoder/block.1.0/mlp/dwconv/dwconv/ReadVariableOp:0"
    output: "tf_segformer_for_semantic_segmentation/segformer/encoder/block.1.0/mlp/dwconv/dwconv/PartitionedCall:0"
    name: "tf_segformer_for_semantic_segmentation/segformer/encoder/block.1.0/mlp/dwconv/dwconv/PartitionedCall"
    op_type: "PartitionedCall"
    ...
    attribute {
      name: "f"
      s: "__inference__jit_compiled_convolution_op_6171"
      type: STRING
    }
  }

The issue is that node __inference__jit_compiled_convolution_op_6171 is referenced, but its definition is nowhere to be found. So likely tf2onnx failed to convert that operation at the first place.

There was a similar issue, where one of tf2onnx contributors said:

StatefulPartitionedCall is an op that does a simple function call in TF. Our converter doesn't normally have to deal with it since the optimizer we run before conversion automatically inlines most function calls. If it shows up in the optimized graph there is usually some reason that will prevent conversion from working.

I created an issue with the details above in tf2onnx GitHub: https://github.com/onnx/tensorflow-onnx/issues/2127

[github-actions[bot]]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.