# Custom model export to onnx-runtime

**URL:** <https://discuss.huggingface.co/t/custom-model-export-to-onnx-runtime/38633>\
**Category:** 🤗Optimum\
**Created:** [May 4, 2023, 8:37am UTC](https://discuss.huggingface.co/t/custom-model-export-to-onnx-runtime/38633 "2023-05-04T08:37:54Z")\
**Posts on this page:** 8\
**Page:** 1

<div class="post-metadata">

**Author:** ![maiia-bocharova](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/maiia-bocharova/32/43850_2.png) [@maiia-bocharova](https://discuss.huggingface.co/u/maiia-bocharova)\
**Post date:** [May 4, 2023, 8:37am UTC](https://discuss.huggingface.co/t/custom-model-export-to-onnx-runtime/38633/1 "2023-05-04T08:37:54Z")

</div>

I modified the BertEmbeddins, BertModel and BertForTokenClassification to accept additional feature (whether token in capitalized or not), in pure transformers it all works, but I am struggling with implementing the export of this custom model (so I can optimize it with optimum and get an inference speed up)

```python
register_for_onnx = TasksManager.create_register("onnx")

@register_for_onnx("custom-bert", "token-classification")
class CustomOnnxConfig(TextEncoderOnnxConfig):
    # Specifies how to normalize the BertConfig, this is needed to access common attributes
    # during dummy input generation.
    NORMALIZED_CONFIG_CLASS = NormalizedTextConfig
    # Sets the absolute tolerance to when validating the exported ONNX model against the
    # reference model.
    ATOL_FOR_VALIDATION = 1e-4

    def inputs(self) -> Dict[str, Dict[int, str]]:
        if self.task == "multiple-choice":
            dynamic_axis = {0: "batch_size", 1: "num_choices", 2: "sequence_length"}
        else:
            dynamic_axis = {0: "batch_size", 1: "sequence_length"}
        return {
            "input_ids": dynamic_axis,
            "attention_mask": dynamic_axis,
            "capitalization_ids": dynamic_axis,
            "token_type_ids": dynamic_axis,
        }

base_model = CustomBertForTokenClassification.from_pretrained("my-checkpoint")

onnx_path = Path("model.onnx")

```

Here I do not understand what to do next  
`base_model.config` returns BertConfig, which I think I need to overwrite with the custom config I created in the previous step.

Can you please help me?

---

<div class="post-metadata">

**Author:** ![michaelbenayoun](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/michaelbenayoun/32/15861_2.png) [@michaelbenayoun](https://discuss.huggingface.co/u/michaelbenayoun)\
**Post date:** [May 4, 2023, 3:45pm UTC](https://discuss.huggingface.co/t/custom-model-export-to-onnx-runtime/38633/2 "2023-05-04T15:45:36Z")

</div>

First, I can see that your new model will have a new input. We need a `DummyInputGenerator` that can handle this. So you could try something like:

```python
from optimum.utils import DummyTextInputGenerator

class MyDummyTextInputenerator(DummyTextInputGenerator):
    SUPPORTED_INPUT_NAMES = (                                                                                                                                                               
        "input_ids",                                                                                                                                                                        
        "attention_mask",                                                                                                                                                                   
        "token_type_ids",
        "capitalization_ids"
      )

class CustomOnnxConfig(TextEncoderOnnxConfig):
    # Specifies how to normalize the BertConfig, this is needed to access common attributes
    # during dummy input generation.
    NORMALIZED_CONFIG_CLASS = NormalizedTextConfig
    DUMMY_INPUT_GENERATOR_CLASSES = (MyDummyTextInputenerator,)
    # Sets the absolute tolerance to when validating the exported ONNX model against the
    # reference model.
    ATOL_FOR_VALIDATION = 1e-4

    @property
    def inputs(self) -> Dict[str, Dict[int, str]]:
        if self.task == "multiple-choice":
            dynamic_axis = {0: "batch_size", 1: "num_choices", 2: "sequence_length"}
        else:
            dynamic_axis = {0: "batch_size", 1: "sequence_length"}
        return {
            "input_ids": dynamic_axis,
            "attention_mask": dynamic_axis,
            "capitalization_ids": dynamic_axis,
            "token_type_ids": dynamic_axis,
        }

```

Since an `OnnxConfig` already exists for `bert`, the register method will not do anything.  
If you want to be able to overwrite existing register configurations you can do that:

```python
register_for_onnx = TasksManager.create_register("onnx", overwrite_existing=True)

```

Another approach is to do it programmatically: you can first export your model using Python code and then optimize it using the CLI.

```python
from pathlib import Path
from optimum.exporters.onnx import export

export(
    model,
    CustomOnnxConfig(model.config),
    Path("my_custom_model_onnx"),
)

```

---

<div class="post-metadata">

**Author:** ![maiia-bocharova](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/maiia-bocharova/32/43850_2.png) [@maiia-bocharova](https://discuss.huggingface.co/u/maiia-bocharova)\
**Post date:** [May 4, 2023, 4:53pm UTC](https://discuss.huggingface.co/t/custom-model-export-to-onnx-runtime/38633/3 "2023-05-04T16:53:30Z")

</div>

Thank you a lot! I am sorry, I still have issues

```python
from optimum.utils import DummyTextInputGenerator
from optimum.exporters.onnx.config import OnnxConfig

class MyDummyTextInputenerator(DummyTextInputGenerator):
    SUPPORTED_INPUT_NAMES = (                                                                                                                                                               
        "input_ids",                                                                                                                                                                        
        "attention_mask",                                                                                                                                                                   
        "token_type_ids",
        "capitalization_ids"
      )
    
class TextEncoderOnnxConfig(OnnxConfig):
    # Describes how to generate the dummy inputs.
    DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator,)

register_for_onnx = TasksManager.create_register("onnx", overwrite_existing=True)

@register_for_onnx("bert", "token-classification")
class CustomOnnxConfig(TextEncoderOnnxConfig):
    # Specifies how to normalize the BertConfig, this is needed to access common attributes
    # during dummy input generation.
    NORMALIZED_CONFIG_CLASS = NormalizedTextConfig
    DUMMY_INPUT_GENERATOR_CLASSES = (MyDummyTextInputenerator,)
    # Sets the absolute tolerance to when validating the exported ONNX model against the
    # reference model.
    ATOL_FOR_VALIDATION = 1e-4

    @property
    def inputs(self) -> Dict[str, Dict[int, str]]:
        if self.task == "multiple-choice":
            dynamic_axis = {0: "batch_size", 1: "num_choices", 2: "sequence_length"}
        else:
            dynamic_axis = {0: "batch_size", 1: "sequence_length"}
        return {
            "input_ids": dynamic_axis,
            "attention_mask": dynamic_axis,
            "capitalization_ids": dynamic_axis,
            "token_type_ids": dynamic_axis,
        }

model = CustomBertForTokenClassification.from_pretrained("my-checkpoint")

from pathlib import Path
from optimum.exporters.onnx import export

export(
    model,
    CustomOnnxConfig(model.config),
    Path("model.onnx"),
)

```

I get error `RuntimeError: number of output names provided (1) exceeded number of outputs (0)`

---

<div class="post-metadata">

**Author:** ![michaelbenayoun](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/michaelbenayoun/32/15861_2.png) [@michaelbenayoun](https://discuss.huggingface.co/u/michaelbenayoun)\
**Post date:** [May 5, 2023, 9:41am UTC](https://discuss.huggingface.co/t/custom-model-export-to-onnx-runtime/38633/4 "2023-05-05T09:41:20Z")

</div>

I think this might be linked to [this issue](https://github.com/huggingface/optimum/issues/1028).  
Could you try installing optimum from sources and re-run your script?

---

<div class="post-metadata">

**Author:** ![maiia-bocharova](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/maiia-bocharova/32/43850_2.png) [@maiia-bocharova](https://discuss.huggingface.co/u/maiia-bocharova)\
**Post date:** [May 5, 2023, 10:44am UTC](https://discuss.huggingface.co/t/custom-model-export-to-onnx-runtime/38633/5 "2023-05-05T10:44:47Z")

</div>

Did installation like this:

```bash
pip install transformers[onnx]==4.28.1
python -m pip install git+https://github.com/huggingface/optimum.git

```

But the error stayed same `RuntimeError: number of output names provided (1) exceeded number of outputs (0)`

---

<div class="post-metadata">

**Author:** ![michaelbenayoun](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/michaelbenayoun/32/15861_2.png) [@michaelbenayoun](https://discuss.huggingface.co/u/michaelbenayoun)\
**Post date:** [May 9, 2023, 6:35pm UTC](https://discuss.huggingface.co/t/custom-model-export-to-onnx-runtime/38633/6 "2023-05-09T18:35:18Z")

</div>

Hi Maiia,  
I saw your message [here](https://colab.research.google.com/drive/1BPaUAXJZ9P0XGjuP08Nht9qwdiwRjGAc?usp=sharing), thanks for creating a notebook, it really helped me to try things out.

You do not need to register anything since you are not using the `TasksManager` in the end. Creating your `OnnxConfig` and using the `export` function is enough.

But because you are not using the `TasksManager`, you have to instantiate the `OnnxConfig` manually, and **you need to specify the name of the task when doing so, otherwise it will infere the task to be the default one** (`BertModel` and not `BertForTokenClassification`).

So to be able to export your model the fix is actually easy:

```python
export(
    model,
    CustomOnnxConfig(model.config, task="token-classification"),
    Path("model.onnx"),
)

```

---

<div class="post-metadata">

**Author:** ![maiia-bocharova](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/maiia-bocharova/32/43850_2.png) [@maiia-bocharova](https://discuss.huggingface.co/u/maiia-bocharova)\
**Post date:** [May 9, 2023, 7:02pm UTC](https://discuss.huggingface.co/t/custom-model-export-to-onnx-runtime/38633/7 "2023-05-09T19:02:13Z")

</div>

> [@michaelbenayoun](#):
>
> here

Thank you a lot! You helped me so, so much!

Last question - if I need to reload it (as `OrtModelForTokenClassification` to optimize with `ORTOptimizer`) I probably need to save a modified config? (If a save original config and do

```python
reloaded_model = ORTModelForTokenClassification.from_pretrained("onnx")
inputs = {k: torch.zeros([2, 16], dtype=torch.long) for k in ["input_ids", "attention_mask", "capitalization_ids", "token_type_ids"]}
reloaded_model(**inputs)

```

I get an error about `ValueError: Model requires 4 inputs. Input Feed contains 3`

Or maybe I just need to modify `OrtModelForTokenClassification` class and it will fix the issue…

---

<div class="post-metadata">

**Author:** ![maiia-bocharova](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/maiia-bocharova/32/43850_2.png) [@maiia-bocharova](https://discuss.huggingface.co/u/maiia-bocharova)\
**Post date:** [May 9, 2023, 7:16pm UTC](https://discuss.huggingface.co/t/custom-model-export-to-onnx-runtime/38633/8 "2023-05-09T19:16:34Z")

</div>

Modified `OrtForTokenClassification` - and everything works perfectly now!

Thank you so, so, so much for patience and kindness and spending time on a beginner!

And for wonderful work which Huggingface gives people! Thank you!!!