Add ONNX export support for DETA model

docs/source/exporters/onnx/overview.mdx

@@ -37,6 +37,7 @@ Supported architectures from [🤗 Transformers](https://huggingface.co/docs/tra
 - Deberta
 - Deberta-v2
 - Deit
+- Deta
 - Detr
 - DistilBert
 - Donut-Swin

optimum/exporters/onnx/model_configs.py

@@ -790,6 +790,56 @@ class ResNetOnnxConfig(ViTOnnxConfig):
     ATOL_FOR_VALIDATION = 1e-3
     DEFAULT_ONNX_OPSET = 11
 
+class DetaDummyInputGenerator(DummyVisionInputGenerator):
+    def __init__(
+        self,
+        task: str,
+        normalized_config: NormalizedVisionConfig,
+        batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"],
+        num_channels: int = DEFAULT_DUMMY_SHAPES["num_channels"],
+        width: int = DEFAULT_DUMMY_SHAPES["width"],
+        height: int = DEFAULT_DUMMY_SHAPES["height"],
+        **kwargs,
+    ):
+        super().__init__(
+            task=task,
+            normalized_config=normalized_config,
+            batch_size=batch_size,
+            num_channels=num_channels,
+            width=width,
+            height=height,
+            **kwargs,
+        )
+
+        # from transformers.onnx.utils import get_preprocessor
+
+        # preprocessor = get_preprocessor(normalized_config._name_or_path)
+        # if preprocessor is not None and hasattr(preprocessor, "crop_size"):
+        #     self.height = preprocessor.crop_size.get("height", self.height)
+        #     self.width = preprocessor.crop_size.get("width", self.width)
+        self.width = 1066
+        self.height = 800
+
+    def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"):
+        input_ = super().generate(
+            input_name=input_name, framework=framework, int_dtype=int_dtype, float_dtype=float_dtype
+        )
+        return input_
+
+class DetaOnnxConfig(ViTOnnxConfig):
+    # OPSET=16 required. Otherwise we get the following error:
+    # torch.onnx.errors.UnsupportedOperatorError: Exporting the operator 'aten::grid_sampler' to ONNX opset version 12 is not supported. Support for this operator was added in version 16, try exporting with this version.
+    DEFAULT_ONNX_OPSET = 16
+    DUMMY_INPUT_GENERATOR_CLASSES = (DetaDummyInputGenerator, )
+    ATOL_FOR_VALIDATION = 1e-3
+
+    @property
+    def inputs(self) -> Dict[str, Dict[int, str]]:
+        return {
+            "pixel_values": {0: "batch_size"},
+            'pixel_mask': {0: "batch_size"}
+               }
+
 
 class DetrOnnxConfig(ViTOnnxConfig):
     DEFAULT_ONNX_OPSET = 12

optimum/exporters/tasks.py

@@ -568,6 +568,11 @@ class TasksManager:
             "masked-im",
             onnx="DeiTOnnxConfig",
         ),
+        "deta": supported_tasks_mapping(
+            "feature-extraction",
+            "object-detection",
+            onnx="DetaOnnxConfig"
+        ),
         "detr": supported_tasks_mapping(
             "feature-extraction",
             "object-detection",

optimum/utils/normalized_config.py

@@ -203,6 +203,7 @@ class NormalizedConfigManager:
         'convnextv2',
         'data2vec-text',
         'data2vec-vision',
+        'deta',
         'detr',
         'flaubert',
         'groupvit',

tests/exporters/exporters_utils.py

@@ -220,6 +220,7 @@
     "deberta": "hf-internal-testing/tiny-random-DebertaModel",  # Not using microsoft/deberta-base because it takes too much time for testing.
     "deberta-v2": "hf-internal-testing/tiny-random-DebertaV2Model",  # Not using microsoft/deberta-v2-xlarge because it takes too much time for testing.
     "deit": "facebook/deit-small-patch16-224",
+    "deta": "jozhang97/deta-resnet-50",
     "detr": "hf-internal-testing/tiny-random-detr",  # Not using facebook/detr-resnet-50 because it takes too much time for testing.
     "distilbert": "distilbert-base-cased",
     "electra": "google/electra-base-generator",