Invalid graph when converting TFBertForSequenceClassification

[ZJaume]

Are text classifiaction huggingface models supported? I'm getting invalid graph when running this:

import os
from transformers import TFBertForSequenceClassification
import keras2onnx
model_keras = TFBertForSequenceClassification.from_pretrained('bert-base-cased')
model_onnx = keras2onnx.convert_keras(model_keras, model_keras.name)

import psutil
import onnxruntime
os.environ["OMP_NUM_THREADS"] = str(psutil.cpu_count(logical=True))
os.environ["OMP_WAIT_POLICY"] = 'ACTIVE'

sess_options = onnxruntime.SessionOptions()
session = onnxruntime.InferenceSession(model_onnx.SerializeToString(), sess_options)

All model checkpoint layers were used when initializing TFBertForSequenceClassification.

Some layers of TFBertForSequenceClassification were not initialized from the model checkpoint at bert-base-cased and are newly initialized: ['classifier']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
tf executing eager_mode: True
tf.keras model eager_mode: False
WARNING:tensorflow:The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).
WARNING:tensorflow:The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.
The ONNX operator number change on the optimization: 2428 -> 1618
Traceback (most recent call last):
  File "convert_example.py", line 16, in <module>
    session = onnxruntime.InferenceSession(model_onnx.SerializeToString(), sess_options)
  File "/work/jzaragoza/bicleaner-neural/venv/lib/python3.7/site-packages/onnxruntime/capi/onnxruntime_inference_collection.py", line 280, in __init__
    self._create_inference_session(providers, provider_options)
  File "/work/jzaragoza/bicleaner-neural/venv/lib/python3.7/site-packages/onnxruntime/capi/onnxruntime_inference_collection.py", line 309, in _create_inference_session
    sess = C.InferenceSession(session_options, self._model_bytes, False, self._read_config_from_model)
onnxruntime.capi.onnxruntime_pybind11_state.InvalidGraph: [ONNXRuntimeError] : 10 : INVALID_GRAPH : This is an invalid model. Error in Node:tf_bert_for_sequence_classification/bert/Shape:0_squeeze : Unrecognized attribute: axes for operator Squeeze

[ZJaume]

Sorry, forgot the versions.

keras2onnx==1.7.0
onnx==1.9.0
onnxconverter-common==1.8.1
onnxruntime==1.7.0
tensorflow==2.3.2
transformers==4.4.2

[TruscaPetre]

Try this converter, it has worked for me: https://github.com/onnx/onnxmltools . Although you should take into account that I used tensorflow 2.2 for training the model, so make this adjustment also.

[ZJaume]

Yeah, I tried tf2onnx (which seems to be wrapped in that tool) and it worked.