Repository: huggingface/transformers
Issue category: Feature Request

Enable ONNX export for `VisionDecoderEncoderModel`
# 🚀 Feature request

<!-- A clear and concise description of the feature proposal.
     Please provide a link to the paper and code in case they exist. -->

As discussed on the [forums](https://discuss.huggingface.co/t/can-i-export-a-visionencoderdecoder-checkpoint-to-onnx/12885), it would be nice if one could export `VisionDecoderEncoderModel` classes using the `transformers.onnx` package.

## Motivation

<!-- Please outline the motivation for the proposal. Is your feature request
     related to a problem? e.g., I'm always frustrated when [...]. If this is related
     to another GitHub issue, please link here too. -->

It is currently not possible to export `VisionDecoderEncoderModel` classes (or vision models more generally), unless the end user is willing to write their own `OnnxConfig`. It might make sense to first explore what is involved in exporting ViT before tackling this more complex example. 

Looking at this a bit more closely, I can see that `transformers.onnx` currently has a tight integration with tokenizers (e.g. to generate dummy inputs), so some refactoring will be necessary to support other modalities.



Functions to edit:
src/transformers/onnx/__main__.py:main
src/transformers/onnx/config.py:OnnxConfig
src/transformers/onnx/features.py:FeaturesManager
Functions added: src/transformers/onnx/features.py:FeaturesManager

Patch:
diff --git a/docs/source/en/serialization.mdx b/docs/source/en/serialization.mdx
index a1577447f723..95df1d883a26 100644
--- a/docs/source/en/serialization.mdx
+++ b/docs/source/en/serialization.mdx
@@ -95,6 +95,7 @@ Ready-made configurations include the following architectures:
 - SegFormer
 - SqueezeBERT
 - T5
+- Vision Encoder decoder
 - ViT
 - XLM
 - XLM-RoBERTa
@@ -298,6 +299,13 @@ in the attention blocks) that can be used for fast autoregressive decoding.
 
 </Tip>
 
+<Tip>
+
+For `VisionEncoderDecoder` type models, the encoder and decoder parts are
+exported separately as two ONNX files named `encoder_model.onnx` and `decoder_model.onnx` respectively.
+
+</Tip>
+
 
 ### Exporting a model for an unsupported architecture
 
diff --git a/src/transformers/models/vision_encoder_decoder/__init__.py b/src/transformers/models/vision_encoder_decoder/__init__.py
index 5d501b8feb83..fcb53d9d1337 100644
--- a/src/transformers/models/vision_encoder_decoder/__init__.py
+++ b/src/transformers/models/vision_encoder_decoder/__init__.py
@@ -27,7 +27,9 @@
 )
 
 
-_import_structure = {"configuration_vision_encoder_decoder": ["VisionEncoderDecoderConfig"]}
+_import_structure = {
+    "configuration_vision_encoder_decoder": ["VisionEncoderDecoderConfig", "VisionEncoderDecoderOnnxConfig"]
+}
 
 try:
     if not is_torch_available():
@@ -54,7 +56,7 @@
     _import_structure["modeling_flax_vision_encoder_decoder"] = ["FlaxVisionEncoderDecoderModel"]
 
 if TYPE_CHECKING:
-    from .configuration_vision_encoder_decoder import VisionEncoderDecoderConfig
+    from .configuration_vision_encoder_decoder import VisionEncoderDecoderConfig, VisionEncoderDecoderOnnxConfig
 
     try:
         if not is_torch_available():
diff --git a/src/transformers/models/vision_encoder_decoder/configuration_vision_encoder_decoder.py b/src/transformers/models/vision_encoder_decoder/configuration_vision_encoder_decoder.py
index b2c3b2aaccaa..14ada4c8b7bf 100644
--- a/src/transformers/models/vision_encoder_decoder/configuration_vision_encoder_decoder.py
+++ b/src/transformers/models/vision_encoder_decoder/configuration_vision_encoder_decoder.py
@@ -15,12 +15,19 @@
 # limitations under the License.
 
 import copy
+from typing import TYPE_CHECKING, Any, Mapping, Optional, OrderedDict
+
+from packaging import version
 
 from ...configuration_utils import PretrainedConfig
+from ...onnx import OnnxConfig
 from ...utils import logging
 from ..auto.configuration_auto import AutoConfig
 
 
+if TYPE_CHECKING:
+    from ... import PreTrainedTokenizerBase, TensorType
+
 logger = logging.get_logger(__name__)
 
 
@@ -119,3 +126,97 @@ def to_dict(self):
         output["decoder"] = self.decoder.to_dict()
         output["model_type"] = self.__class__.model_type
         return output
+
+
+class VisionEncoderDecoderEncoderOnnxConfig(OnnxConfig):
+    torch_onnx_minimum_version = version.parse("1.11")
+
+    @property
+    def inputs(self) -> Mapping[str, Mapping[int, str]]:
+        return OrderedDict(
+            [
+                ("pixel_values", {0: "batch", 1: "num_channels", 2: "height", 3: "width"}),
+            ]
+        )
+
+    @property
+    def atol_for_validation(self) -> float:
+        return 1e-4
+
+    @property
+    def outputs(self) -> Mapping[str, Mapping[int, str]]:
+        return OrderedDict({"last_hidden_state": {0: "batch", 1: "encoder_sequence"}})
+
+
+class VisionEncoderDecoderDecoderOnnxConfig(OnnxConfig):
+    @property
+    def inputs(self) -> Mapping[str, Mapping[int, str]]:
+        common_inputs = OrderedDict()
+        common_inputs["input_ids"] = {0: "batch", 1: "past_decoder_sequence + sequence"}
+        common_inputs["attention_mask"] = {0: "batch", 1: "past_decoder_sequence + sequence"}
+        common_inputs["encoder_hidden_states"] = {0: "batch", 1: "encoder_sequence"}
+
+        return common_inputs
+
+    def generate_dummy_inputs(
+        self,
+        tokenizer: "PreTrainedTokenizerBase",
+        batch_size: int = -1,
+        seq_length: int = -1,
+        is_pair: bool = False,
+        framework: Optional["TensorType"] = None,
+    ) -> Mapping[str, Any]:
+        import torch
+
+        common_inputs = OrderedDict()
+
+        dummy_input = super().generate_dummy_inputs(
+            tokenizer, batch_size=batch_size, seq_length=seq_length, is_pair=is_pair, framework=framework
+        )
+
+        batch, encoder_sequence = dummy_input["input_ids"].shape
+        encoder_hidden_states_shape = (batch, encoder_sequence, self._config.encoder_hidden_size)
+        common_inputs["input_ids"] = dummy_input.pop("input_ids")
+        common_inputs["attention_mask"] = dummy_input.pop("attention_mask")
+        common_inputs["encoder_hidden_states"] = torch.zeros(encoder_hidden_states_shape)
+
+        return common_inputs
+
+
+class VisionEncoderDecoderOnnxConfig(OnnxConfig):
+    @property
+    def inputs(self) -> None:
+        pass
+
+    def get_encoder_config(self, encoder_config: PretrainedConfig) -> OnnxConfig:
+        r"""
+        Returns ONNX encoder config for `VisionEncoderDecoder` model.
+
+        Args:
+            encoder_config (`PretrainedConfig`):
+                The encoder model's configuration to use when exporting to ONNX.
+
+        Returns:
+            [`VisionEncoderDecoderEncoderOnnxConfig`]: An instance of the ONNX configuration object
+        """
+        return VisionEncoderDecoderEncoderOnnxConfig(encoder_config)
+
+    def get_decoder_config(
+        self, encoder_config: PretrainedConfig, decoder_config: PretrainedConfig, feature: str = "default"
+    ) -> OnnxConfig:
+        r"""
+        Returns ONNX decoder config for `VisionEncoderDecoder` model.
+
+        Args:
+            encoder_config (`PretrainedConfig`):
+                The encoder model's configuration to use when exporting to ONNX.
+            decoder_config (`PretrainedConfig`):
+                The decoder model's configuration to use when exporting to ONNX
+            feature (`str`, *optional*):
+                The type of feature to export the model with.
+
+        Returns:
+            [`VisionEncoderDecoderDecoderOnnxConfig`]: An instance of the ONNX configuration object.
+        """
+        decoder_config.encoder_hidden_size = encoder_config.hidden_size
+        return VisionEncoderDecoderDecoderOnnxConfig(decoder_config, feature)
diff --git a/src/transformers/onnx/__main__.py b/src/transformers/onnx/__main__.py
index 55ad5f54c994..b84e12edbb24 100644
--- a/src/transformers/onnx/__main__.py
+++ b/src/transformers/onnx/__main__.py
@@ -22,6 +22,9 @@
 from .features import FeaturesManager
 
 
+ENCODER_DECODER_MODELS = ["vision-encoder-decoder"]
+
+
 def main():
     parser = ArgumentParser("Hugging Face Transformers ONNX exporter")
     parser.add_argument(
@@ -65,48 +68,110 @@ def main():
     if not args.output.parent.exists():
         args.output.parent.mkdir(parents=True)
 
-    # Instantiate the appropriate preprocessor
-    if args.preprocessor == "auto":
-        preprocessor = get_preprocessor(args.model)
-    elif args.preprocessor == "tokenizer":
-        preprocessor = AutoTokenizer.from_pretrained(args.model)
-    elif args.preprocessor == "feature_extractor":
-        preprocessor = AutoFeatureExtractor.from_pretrained(args.model)
-    elif args.preprocessor == "processor":
-        preprocessor = AutoProcessor.from_pretrained(args.model)
-    else:
-        raise ValueError(f"Unknown preprocessor type '{args.preprocessor}'")
-
     # Allocate the model
     model = FeaturesManager.get_model_from_feature(
         args.feature, args.model, framework=args.framework, cache_dir=args.cache_dir
     )
+
     model_kind, model_onnx_config = FeaturesManager.check_supported_model_or_raise(model, feature=args.feature)
     onnx_config = model_onnx_config(model.config)
 
-    # Ensure the requested opset is sufficient
-    if args.opset is None:
-        args.opset = onnx_config.default_onnx_opset
+    if model_kind in ENCODER_DECODER_MODELS:
+        encoder_model = model.get_encoder()
+        decoder_model = model.get_decoder()
 
-    if args.opset < onnx_config.default_onnx_opset:
-        raise ValueError(
-            f"Opset {args.opset} is not sufficient to export {model_kind}. "
-            f"At least  {onnx_config.default_onnx_opset} is required."
+        encoder_onnx_config = onnx_config.get_encoder_config(encoder_model.config)
+        decoder_onnx_config = onnx_config.get_decoder_config(
+            encoder_model.config, decoder_model.config, feature=args.feature
         )
 
-    onnx_inputs, onnx_outputs = export(
-        preprocessor,
-        model,
-        onnx_config,
-        args.opset,
-        args.output,
-    )
+        if args.opset is None:
+            args.opset = max(encoder_onnx_config.default_onnx_opset, decoder_onnx_config.default_onnx_opset)
+
+        if args.opset < min(encoder_onnx_config.default_onnx_opset, decoder_onnx_config.default_onnx_opset):
+            raise ValueError(
+                f"Opset {args.opset} is not sufficient to export {model_kind}. At least "
+                f" {min(encoder_onnx_config.default_onnx_opset, decoder_onnx_config.default_onnx_opset)} is required."
+            )
+
+        preprocessor = AutoFeatureExtractor.from_pretrained(args.model)
+
+        onnx_inputs, onnx_outputs = export(
+            preprocessor,
+            encoder_model,
+            encoder_onnx_config,
+            args.opset,
+            args.output.parent.joinpath("encoder_model.onnx"),
+        )
+
+        validate_model_outputs(
+            encoder_onnx_config,
+            preprocessor,
+            encoder_model,
+            args.output.parent.joinpath("encoder_model.onnx"),
+            onnx_outputs,
+            args.atol if args.atol else encoder_onnx_config.atol_for_validation,
+        )
+
+        preprocessor = AutoTokenizer.from_pretrained(args.model)
+
+        onnx_inputs, onnx_outputs = export(
+            preprocessor,
+            decoder_model,
+            decoder_onnx_config,
+            args.opset,
+            args.output.parent.joinpath("decoder_model.onnx"),
+        )
+
+        validate_model_outputs(
+            decoder_onnx_config,
+            preprocessor,
+            decoder_model,
+            args.output.parent.joinpath("decoder_model.onnx"),
+            onnx_outputs,
+            args.atol if args.atol else decoder_onnx_config.atol_for_validation,
+        )
+        logger.info(
+            f"All good, model saved at: {args.output.parent.joinpath('encoder_model.onnx').as_posix()},"
+            f" {args.output.parent.joinpath('decoder_model.onnx').as_posix()}"
+        )
+
+    else:
+        # Instantiate the appropriate preprocessor
+        if args.preprocessor == "auto":
+            preprocessor = get_preprocessor(args.model)
+        elif args.preprocessor == "tokenizer":
+            preprocessor = AutoTokenizer.from_pretrained(args.model)
+        elif args.preprocessor == "feature_extractor":
+            preprocessor = AutoFeatureExtractor.from_pretrained(args.model)
+        elif args.preprocessor == "processor":
+            preprocessor = AutoProcessor.from_pretrained(args.model)
+        else:
+            raise ValueError(f"Unknown preprocessor type '{args.preprocessor}'")
+
+        # Ensure the requested opset is sufficient
+        if args.opset is None:
+            args.opset = onnx_config.default_onnx_opset
+
+        if args.opset < onnx_config.default_onnx_opset:
+            raise ValueError(
+                f"Opset {args.opset} is not sufficient to export {model_kind}. "
+                f"At least  {onnx_config.default_onnx_opset} is required."
+            )
+
+        onnx_inputs, onnx_outputs = export(
+            preprocessor,
+            model,
+            onnx_config,
+            args.opset,
+            args.output,
+        )
 
-    if args.atol is None:
-        args.atol = onnx_config.atol_for_validation
+        if args.atol is None:
+            args.atol = onnx_config.atol_for_validation
 
-    validate_model_outputs(onnx_config, preprocessor, model, args.output, onnx_outputs, args.atol)
-    logger.info(f"All good, model saved at: {args.output.as_posix()}")
+        validate_model_outputs(onnx_config, preprocessor, model, args.output, onnx_outputs, args.atol)
+        logger.info(f"All good, model saved at: {args.output.as_posix()}")
 
 
 if __name__ == "__main__":
diff --git a/src/transformers/onnx/config.py b/src/transformers/onnx/config.py
index 3b789051a220..dc27c0b6924d 100644
--- a/src/transformers/onnx/config.py
+++ b/src/transformers/onnx/config.py
@@ -103,6 +103,7 @@ class OnnxConfig(ABC):
         "seq2seq-lm": OrderedDict({"logits": {0: "batch", 1: "decoder_sequence"}}),
         "sequence-classification": OrderedDict({"logits": {0: "batch"}}),
         "token-classification": OrderedDict({"logits": {0: "batch", 1: "sequence"}}),
+        "vision2seq-lm": OrderedDict({"logits": {0: "batch", 1: "sequence"}}),
     }
 
     def __init__(self, config: "PretrainedConfig", task: str = "default", patching_specs: List[PatchingSpec] = None):
@@ -451,7 +452,6 @@ def generate_dummy_inputs(
         is_pair: bool = False,
         framework: Optional[TensorType] = None,
     ) -> Mapping[str, Any]:
-
         # TODO: should we set seq_length = 1 when self.use_past = True?
         common_inputs = super().generate_dummy_inputs(
             tokenizer, batch_size=batch_size, seq_length=seq_length, is_pair=is_pair, framework=framework
@@ -577,7 +577,6 @@ def generate_dummy_inputs(
         is_pair: bool = False,
         framework: Optional[TensorType] = None,
     ) -> Mapping[str, Any]:
-
         encoder_inputs = super(OnnxConfigWithPast, self).generate_dummy_inputs(
             tokenizer, batch_size=batch_size, seq_length=seq_length, is_pair=is_pair, framework=framework
         )
diff --git a/src/transformers/onnx/features.py b/src/transformers/onnx/features.py
index 535686f179a5..5bf913efe5f7 100644
--- a/src/transformers/onnx/features.py
+++ b/src/transformers/onnx/features.py
@@ -30,6 +30,7 @@
         AutoModelForSeq2SeqLM,
         AutoModelForSequenceClassification,
         AutoModelForTokenClassification,
+        AutoModelForVision2Seq,
     )
 if is_tf_available():
     from transformers.models.auto import (
@@ -98,6 +99,7 @@ class FeaturesManager:
             "image-segmentation": AutoModelForImageSegmentation,
             "masked-im": AutoModelForMaskedImageModeling,
             "semantic-segmentation": AutoModelForSemanticSegmentation,
+            "vision2seq-lm": AutoModelForVision2Seq,
         }
     if is_tf_available():
         _TASKS_TO_TF_AUTOMODELS = {
@@ -478,6 +480,9 @@ class FeaturesManager:
             "seq2seq-lm-with-past",
             onnx_config_cls="models.t5.T5OnnxConfig",
         ),
+        "vision-encoder-decoder": supported_features_mapping(
+            "vision2seq-lm", onnx_config_cls="models.vision_encoder_decoder.VisionEncoderDecoderOnnxConfig"
+        ),
         "vit": supported_features_mapping(
             "default", "image-classification", "masked-im", onnx_config_cls="models.vit.ViTOnnxConfig"
         ),
@@ -579,6 +584,7 @@ def get_model_class_for_feature(feature: str, framework: str = "pt") -> Type:
             raise KeyError(
                 f"Unknown task: {feature}. Possible values are {list(FeaturesManager._TASKS_TO_AUTOMODELS.values())}"
             )
+
         return task_to_automodel[task]
 
     @staticmethod

Test patch: diff --git a/tests/onnx/test_onnx_v2.py b/tests/onnx/test_onnx_v2.py
index f3c19ed8fa98..1b886320ae73 100644
--- a/tests/onnx/test_onnx_v2.py
+++ b/tests/onnx/test_onnx_v2.py
@@ -161,7 +161,6 @@ def test_values_override(self):
         """
         for name, config in OnnxConfigWithPastTestCaseV2.SUPPORTED_WITH_PAST_CONFIGS:
             with self.subTest(name):
-
                 # without past
                 onnx_config_default = OnnxConfigWithPast.from_model_config(config())
                 self.assertIsNotNone(onnx_config_default.values_override, "values_override should not be None")
@@ -219,6 +218,10 @@ def test_values_override(self):
     ("segformer", "nvidia/segformer-b0-finetuned-ade-512-512"),
 }
 
+PYTORCH_EXPORT_ENCODER_DECODER_MODELS = {
+    ("vision-encoder-decoder", "nlpconnect/vit-gpt2-image-captioning"),
+}
+
 PYTORCH_EXPORT_WITH_PAST_MODELS = {
     ("bloom", "bigscience/bloom-560m"),
     ("gpt2", "gpt2"),
@@ -328,6 +331,70 @@ def _onnx_export(self, test_name, name, model_name, feature, onnx_config_class_c
             except (RuntimeError, ValueError) as e:
                 self.fail(f"{name}, {feature} -> {e}")
 
+    def _onnx_export_encoder_decoder_models(
+        self, test_name, name, model_name, feature, onnx_config_class_constructor, device="cpu"
+    ):
+        from transformers import AutoFeatureExtractor, AutoTokenizer
+        from transformers.onnx import export
+
+        model_class = FeaturesManager.get_model_class_for_feature(feature)
+        config = AutoConfig.from_pretrained(model_name)
+        model = model_class.from_config(config)
+
+        onnx_config = onnx_config_class_constructor(model.config)
+
+        if is_torch_available():
+            from transformers.utils import torch_version
+
+            if torch_version < onnx_config.torch_onnx_minimum_version:
+                pytest.skip(
+                    "Skipping due to incompatible PyTorch version. Minimum required is"
+                    f" {onnx_config.torch_onnx_minimum_version}, got: {torch_version}"
+                )
+
+        encoder_model = model.get_encoder()
+        decoder_model = model.get_decoder()
+
+        encoder_onnx_config = onnx_config.get_encoder_config(encoder_model.config)
+        decoder_onnx_config = onnx_config.get_decoder_config(encoder_model.config, decoder_model.config, feature)
+
+        preprocessor = AutoFeatureExtractor.from_pretrained(model_name)
+
+        onnx_opset = max(encoder_onnx_config.default_onnx_opset, decoder_onnx_config.default_onnx_opset)
+
+        with NamedTemporaryFile("w") as encoder_output:
+            onnx_inputs, onnx_outputs = export(
+                preprocessor, encoder_model, encoder_onnx_config, onnx_opset, Path(encoder_output.name), device=device
+            )
+            validate_model_outputs(
+                encoder_onnx_config,
+                preprocessor,
+                encoder_model,
+                Path(encoder_output.name),
+                onnx_outputs,
+                encoder_onnx_config.atol_for_validation,
+            )
+
+        preprocessor = AutoTokenizer.from_pretrained(model_name)
+
+        with NamedTemporaryFile("w") as decoder_output:
+            onnx_inputs, onnx_outputs = export(
+                preprocessor,
+                decoder_model,
+                decoder_onnx_config,
+                onnx_config.default_onnx_opset,
+                Path(decoder_output.name),
+                device=device,
+            )
+            validate_model_outputs(
+                decoder_onnx_config,
+                preprocessor,
+                decoder_model,
+                Path(decoder_output.name),
+                onnx_outputs,
+                decoder_onnx_config.atol_for_validation,
+            )
+
     @parameterized.expand(_get_models_to_test(PYTORCH_EXPORT_MODELS))
     @slow
     @require_torch
@@ -344,6 +411,28 @@ def test_pytorch_export(self, test_name, name, model_name, feature, onnx_config_
     def test_pytorch_export_on_cuda(self, test_name, name, model_name, feature, onnx_config_class_constructor):
         self._onnx_export(test_name, name, model_name, feature, onnx_config_class_constructor, device="cuda")
 
+    @parameterized.expand(_get_models_to_test(PYTORCH_EXPORT_ENCODER_DECODER_MODELS))
+    @slow
+    @require_torch
+    @require_vision
+    @require_rjieba
+    def test_pytorch_export_encoder_decoder_models(
+        self, test_name, name, model_name, feature, onnx_config_class_constructor
+    ):
+        self._onnx_export_encoder_decoder_models(test_name, name, model_name, feature, onnx_config_class_constructor)
+
+    @parameterized.expand(_get_models_to_test(PYTORCH_EXPORT_ENCODER_DECODER_MODELS))
+    @slow
+    @require_torch
+    @require_vision
+    @require_rjieba
+    def test_pytorch_export_encoder_decoder_models_on_cuda(
+        self, test_name, name, model_name, feature, onnx_config_class_constructor
+    ):
+        self._onnx_export_encoder_decoder_models(
+            test_name, name, model_name, feature, onnx_config_class_constructor, device="cuda"
+        )
+
     @parameterized.expand(_get_models_to_test(PYTORCH_EXPORT_WITH_PAST_MODELS))
     @slow
     @require_torch