Add ONNX export support for RT-DETR models

docs/source/exporters/onnx/overview.mdx

@@ -85,6 +85,7 @@ Supported architectures from [🤗 Transformers](https://huggingface.co/docs/tra
 - ResNet
 - Roberta
 - Roformer
+- RT-DETR
 - SAM
 - Segformer
 - SEW

optimum/exporters/onnx/model_configs.py

@@ -787,6 +787,49 @@ def outputs(self) -> Dict[str, Dict[int, str]]:
             return super().outputs
 
 
+class RTDetrDummyInputGenerator(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, "size"):
+            self.height = preprocessor.size.get("height", self.height)
+            self.width = preprocessor.size.get("width", self.width)
+
+    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 RTDetrOnnxConfig(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 = (RTDetrDummyInputGenerator, )
+    ATOL_FOR_VALIDATION = 1e-3
+
+
 class TableTransformerOnnxConfig(DetrOnnxConfig):
     pass
 

optimum/exporters/tasks.py

@@ -951,6 +951,11 @@ class TasksManager:
             onnx="RoFormerOnnxConfig",
             tflite="RoFormerTFLiteConfig",
         ),
+        "rt-detr": supported_tasks_mapping(
+            "feature-extraction",
+            "object-detection",
+            onnx="RTDetrOnnxConfig",
+        ),
         "sam": supported_tasks_mapping(
             "feature-extraction",
             onnx="SamOnnxConfig",

optimum/utils/normalized_config.py

@@ -216,6 +216,7 @@ class NormalizedConfigManager:
         'owlvit',
         'perceiver',
         'roformer',
+        'rt-detr',
         'squeezebert',
         'table-transformer',
     """