Can't reproduce export to onnx with custom bert model

🐛 Bug

I try to run onnx export on a custom bert model, but during inference I get the following error. I share a google colab with the minimum changes to reproduce. All changes are marked with a # CHANGE comment. https://colab.research.google.com/drive/1eiqyQmvhwGih6IHrOg7MkLSc2q0zMHmH?usp=sharing

InvalidArgument: [ONNXRuntimeError] : 2 : INVALID_ARGUMENT : Non-zero status code returned while running Gather node. Name:'Gather_32' Status Message: indices element out of data bounds, idx=1 must be within the inclusive range [-1,0]

Information

Model I am using (Bert, XLNet …): Bert

Language I am using the model on (English, Chinese …): None, I’m using a custom bert model, and for this bug report I’m using a random bert model.

The problem arises when using: The official example notebook: https://github.com/huggingface/transformers/blob/master/notebooks/04-onnx-export.ipynb

To reproduce

Steps to reproduce the behavior:

Run the convert to onxx script with a custom bert model. I’ve made a copy of the official notebook with the minimum changes required to illustrate the problem here: https://colab.research.google.com/drive/1eiqyQmvhwGih6IHrOg7MkLSc2q0zMHmH?usp=sharing

---------------------------------------------------------------------------

InvalidArgument                           Traceback (most recent call last)

<ipython-input-12-1d032f1e9ad0> in <module>()
      9 
     10 # Run the model (None = get all the outputs)
---> 11 sequence, pooled = cpu_model.run(None, inputs_onnx)
     12 
     13 # Print information about outputs

/usr/local/lib/python3.6/dist-packages/onnxruntime/capi/session.py in run(self, output_names, input_feed, run_options)
    109             output_names = [output.name for output in self._outputs_meta]
    110         try:
--> 111             return self._sess.run(output_names, input_feed, run_options)
    112         except C.EPFail as err:
    113             if self._enable_fallback:

InvalidArgument: [ONNXRuntimeError] : 2 : INVALID_ARGUMENT : Non-zero status code returned while running Gather node. Name:'Gather_32' Status Message: indices element out of data bounds, idx=1 must be within the inclusive range [-1,0]

Expected behavior

Get pooled and sequence output of bert model.

Environment info

• transformers version: 2.10.0
• Platform: Linux-4.19.104±x86_64-with-Ubuntu-18.04-bionic
• Python version: 3.6.9
• PyTorch version (GPU?): 1.5.0+cu101 (True)
• Tensorflow version (GPU?): 2.2.0 (True)
• Using GPU in script?: yes
• Using distributed or parallel set-up in script?: no

[mfuntowicz]

@RensDimmendaal I think your suggestion is the way to go, do you mind submitting a PR and assigning me as a reviewer ? 👍

[tianleiwu]

@mfuntowicz,

I run the notebook in my local machine, and look at the onnx model after export (and before optimization). I found that the exported onnx model has switched the position of “attention_mask” with “token_type_ids”:

The above is a snapshot of embedding layer in exported graph. The “attention_mask” in the graph shall be named as “token_type_ids” since it is used to look up segment embeddings.