Documentation for exporting custom architecture to ONNX
Author: ndobb
Pull request: no

I have a custom Bert model for relation classification based on the [R-BERT paper](https://arxiv.org/abs/1905.08284). The model performs well but is relatively slow on CPU, so I'd like to try exporting to ONNX.

The model inherits from `BertPreTrainedModel` and is relatively simple:

```python
class BertForRelationClassification(BertPreTrainedModel):
    def __init__(self, config):
        super(BertForRelationClassification, self).__init__(config)
        self.num_labels = config.num_labels
        self.bert = BertModel(config)
        self.cls_dropout = nn.Dropout(0.1)  
        self.ent_dropout = nn.Dropout(0.1)
        self.classifier = nn.Linear(config.hidden_size*3, self.config.num_labels)
        self.init_weights()

    def forward(self, input_ids, token_type_ids=None, attention_mask=None, e1_mask=None, e2_mask=None, 
                        labels=None, position_ids=None, head_mask=None):
        outputs = self.bert(input_ids, position_ids=position_ids, token_type_ids=token_type_ids,
                            attention_mask=attention_mask, head_mask=head_mask)
        pooled_output = outputs[1]  
        sequence_output = outputs[0] 

        def extract_entity(sequence_output, e_mask):
            extended_e_mask = e_mask.unsqueeze(1)
            extended_e_mask = torch.bmm(
                extended_e_mask.float(), sequence_output).squeeze(1)
            return extended_e_mask.float()

        e1_h = self.ent_dropout(extract_entity(sequence_output, e1_mask))
        e2_h = self.ent_dropout(extract_entity(sequence_output, e2_mask))
        context = self.cls_dropout(pooled_output)
        pooled_output = torch.cat([context, e1_h, e2_h], dim=-1)

        logits = self.classifier(pooled_output)
        outputs = (logits,) + outputs[2:]

        if labels is not None:
            loss_fct = CrossEntropyLoss()
            loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))

        outputs = (loss,) + outputs
        return outputs 
```

Following the [guidelines for exporting custom models to ONNX](https://huggingface.co/transformers/serialization.html#implementing-a-custom-configuration-for-an-unsupported-architecture), I've created a custom OnnxConfig for it and specified inputs and outputs:

```python

class BertForRelationClassificationOnnxConfig(OnnxConfig):
    @property
    def inputs(self) -> Mapping[str, Mapping[int, str]]:
        return OrderedDict(
            [
                ("input_ids", {0: "batch", 1: "sequence"}),
                ("attention_mask", {0: "batch", 1: "sequence"}),
                ("token_type_ids", {0: "batch", 1: "sequence"}),
                ("labels", {0: "batch", 1: "sequence"}),
                ("e1_mask", {0: "batch", 1: "sequence"}),
                ("e2_mask", {0: "batch", 1: "sequence"}),
            ]
        )

    @property
    def outputs(self) -> Mapping[str, Mapping[int, str]]:
        return OrderedDict([("outputs", {0: "batch", 1: "sequence"})]) 
```

However when I run `convert_graph_to_onnx.py`, (of course) the model is assumed to be `BertModel` and the inputs and outputs are those of the vanilla `BertModel`. I'm unclear on the next steps.

I'm fairly sure this is what I should do next as stated on the documentation page: ***"Once this is done, a single step remains: adding this configuration object to the initialisation of the model class, and to the general `transformers` initialisation"***. 

While I feel a bit dense I'm still not following how to make this work, as the `BertForRelationClassificationOnnxConfig` class I created doesn't inherit from `BertConfig` (I could make it do so, but the documentation doesn't specify this) so I don't see how I can use this for initialization of the model. The [MBart example](https://github.com/huggingface/transformers/pull/13049/commits/d097adcebd89a520f04352eb215a85916934204f) doesn't make sense to me as I'm not contributing to the transformers code base. 

Can you please provide guidance or a specific example? Thank you!

Comments:
Hi this is what I did for `CamemBERT` TokenClassification model:
- git clone https://github.com/huggingface/transformers.git
- cd transformers
- open code editor and find the file `src/transformers/onnx/features.py`
- add a new item to `_SUPPORTED_MODEL_KIND` dictionnary

Because my camembert model have the same config, I used the `RobertaOnnxConfig`
```diff
"roberta": supported_features_mapping("default", onnx_config_cls=RobertaOnnxConfig),
+ "camembert": supported_features_mapping("default", onnx_config_cls=RobertaOnnxConfig),
```

- Rebuild the transformers package with `pip install .` at the root directory
- Relaunch your conversion script

If you don't find a model already added with the same config as your model, then you have to add it to your model configuration file.
Example: for `Roberta` the path is `transformers/src/transformers/models/roberta/configuration_roberta.py` and there is already `RobertaOnnxConfig` class. After that you will have to import it in `features.py` 

```python
from ..models.roberta import RobertaOnnxConfig
```
And add it to `_SUPPORTED_MODEL_KIND` dictionnary and do the rebuild of the transformers package.

I think you have to add your created class to the bert configuration file and import it in features.py file and add it to the dictionnary. It's probably the easiest thing to do in your case.
Tell me if this works for you 🤗

P.S.: The best way is to add it to your model repo and make a PR to add it to the base repository package. This is what I'm going to do for camemBERT as well.

@ChainYo, thank you very much! That makes a lot of sense and I was clearly misunderstanding.

I'll do as you suggest and, if it works, close the issue. 

Thank you @ChainYo for the insightful comment - if the documentation is lacking in any way, we're very open to PRs that would help make it clearer. Thank you both!

> if the documentation is lacking in any way, we're very open to PRs that would help make it clearer. 

Yes, I am going to prepare one 🤗
I am just wondering if this is usefull to add the camemBERT config code to camemBERT model because the configuration is exactly the same as Roberta, it will make duplicates in code in differents files. but adding `"camembert": supported_features_mapping("default", onnx_config_cls=RobertaOnnxConfig),` doesn't make sense in the official transformers repository, it's more a personnal trick. So I don't know what to do! @LysandreJik 

Having a duplicate for CamemBERT isn't an issue :)

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](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.