Can't reproduce export to onnx with custom bert model
Author: RensDimmendaal
Pull request: no
Labels: name: External

# 🐛 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

```python
---------------------------------------------------------------------------

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


Comments:
Pinging @mfuntowicz, chief onnx officer

Hi @RensDimmendaal,

Thanks for reporting this 👍.
Can you share the shape of the input you're feeding to the ONNX model? 

Thanks for investigating!

```python
for k,v in inputs_onnx.items():
  print(f"{k}: shape: {v.shape}")
```
```
>>>
input_ids: shape: (1, 10)
token_type_ids: shape: (1, 10)
attention_mask: shape: (1, 10)
```

Source: https://colab.research.google.com/drive/1eiqyQmvhwGih6IHrOg7MkLSc2q0zMHmH#scrollTo=64GodG5fKb0m&line=2&uniqifier=1

Interesting I'm not able to reproduce on my side (_see at the end_). 

Can you try restarting the Colab kernel? _(just to make sure the path are correctly updated)_. 
Let us know if it change something, if not I'll dig further on a custom colab.

```python
>>> import onnxruntime as ort
>>> from transformers import BertTokenizerFast
>>> session = ort.InferenceSession("onnx/bert-base-cased.onnx")

>>> tokenizer = BertTokenizerFast.from_pretrained("bert-base-cased")
>>> onnx_in = tokenizer.encode_plus("S E Q W E N C E", return_tensors="pt")

>>> inputs_onnx = {k: v.cpu().detach().numpy() for k, v in onnx_in.items()}
>>> sequence, pooled = session.run(None, inputs_onnx)

>>> sequence.shape
(1, 10, 768)

>>> pooled.shape
(1, 768)
```

I've done a restart and run all and the problem persists.

i ran your code too, and it gives the following error:

```
---------------------------------------------------------------------------

Fail                                      Traceback (most recent call last)

<ipython-input-13-6776c93b3fb0> in <module>()
     18 
     19 inputs_onnx = {k: v.cpu().detach().numpy() for k, v in onnx_in.items()}
---> 20 sequence, pooled = session.run(None, inputs_onnx)
     21 
     22 print(sequence.shape, pooled.shape)

/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:

Fail: [ONNXRuntimeError] : 1 : FAIL : Non-zero status code returned while running Attention node. Name:'Attention_1' Status Message: CUBLAS error executing cublasGemmHelper( cublas, CUBLAS_OP_N, CUBLAS_OP_N, n, m, 1, &one, reinterpret_cast<const CudaT*>(bias->template Data<T>()), n, GetConstOnes<CudaT>(m), 1, &zero, reinterpret_cast<CudaT*>(gemm_buffer.get()), n, device_prop)
```

source: https://colab.research.google.com/drive/1eiqyQmvhwGih6IHrOg7MkLSc2q0zMHmH#scrollTo=dXeNg37RTxl_&line=9&uniqifier=1

Ok , thanks for checking @RensDimmendaal .

I'll do some experiments on a fresh notebook and post update here 👍 

@mfuntowicz,

I run the [notebook](https://colab.research.google.com/drive/1eiqyQmvhwGih6IHrOg7MkLSc2q0zMHmH#scrollTo=64GodG5fKb0m&line=2&uniqifier=1) 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":

![image](https://user-images.githubusercontent.com/30328909/82973767-20ef9a00-9f8d-11ea-93d1-0876a65ef824.png)

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.

Hi,

how can I use it for my trained text classifier?

The onnx export script has assumption of order of inputs. If the class you used does not have same order (or there are other parameters in between), you can wrap a class to use the expected order for export like:
```
class MyBertModel(BertForMaskedLM):
    def __init__(self, config):
        super().__init__(config)

    def forward(self, input_ids, token_type_ids, attention_mask):
        return super().forward(input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids)

model = MyBertModel(config)
model.save_pretrained("./my_bert")
```
In this way, the exported model will have correct inputs.

Thanks for checking tianleiwu! 

(quick question in between: how do you make that plot of the onnx exported model?)

It does not solve the issue for me though. The error message remains the same.

I've added your code here:
https://colab.research.google.com/drive/1eiqyQmvhwGih6IHrOg7MkLSc2q0zMHmH#scrollTo=gXQ_JorGpAdI&line=1&uniqifier=1

However, changing this in my inputs during inference did the trick:

```python
# CHANGE: SHUFFLE INPUTS
inputs_onnx = {
    'input_ids': inputs_onnx['input_ids'],
    'attention_mask': inputs_onnx['token_type_ids'],
    'token_type_ids': inputs_onnx['attention_mask'],
 }
# Run the model (None = get all the outputs)
sequence, pooled = cpu_model.run(None, inputs_onnx)
```

```
>>>
Sequence output: (1, 10, 768), Pooled output: (1, 768)
```

To me this seems like a bug that could be solved by having `transformers.convert_graph_to_onnx.ensure_valid_input` also return the reordered input_names.
Something like this:

```python
def ensure_valid_input(model, tokens, input_names):
    """
    Ensure input are presented in the correct order, without any None
    Args:
        model: The model used to forward the input data
        tokens: BatchEncoding holding the input data
        input_names: The name of the inputs

    Returns: Tuple

    """
    model_args_name = model.forward.__code__.co_varnames
    model_args_pos = [(model_args_name.index(name) - 1, name) for name in input_names]
    model_args = [None] * (max(map(lambda x: x[0], model_args_pos)) + 1)
    ordered_input_names = [None] * len(model_args) # new

    for arg_pos, arg_name in model_args_pos:
        model_args[arg_pos] = tokens[arg_name]
        ordered_input_names[arg_pos] = arg_name # new

    model_args = tuple(takewhile(lambda arg: arg is not None, model_args))  # Need to be ordered
    return ordered_input_names, model_args # new
```

However, based on the test for this function it seems that it is also used for GPT2, and I don't know if this change will break anythinig for that model (test_onnx.py line 111 and 112).

Happy to submit a PR if this seeems like the way to go.



@tianleiwu @RensDimmendaal I'll have a look on this asap. 

The export should not permute inputs.

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

Thanks @RensDimmendaal for submitting the PR, I'm closing this for now 👍.

Don't hesitate to reopen / create a new issue if you ran into any problem!