Add ONNX config for RT-DETR (and RT-DETRv2) (#2201)

* Add RTDetr Onnx Config * add test and comment * fix style * Update tests large model * fix expected configs model type * Fix inputs shape and remove pixel_mask * Add rt-detr-v2 * Add tests * Add separate V2 config * Update min size according to num_queries --------- Co-authored-by: Yann HALLOUARD <[email protected]> Co-authored-by: Ella Charlaix <[email protected]>
huggingface · Mar 3, 2025 · b6c2b5c · b6c2b5c
1 parent deec884
commit b6c2b5c
Show file tree

Hide file tree

Showing 3 changed files with 47 additions and 1 deletion.
diff --git a/optimum/exporters/onnx/model_configs.py b/optimum/exporters/onnx/model_configs.py
@@ -13,8 +13,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Model specific ONNX configurations."""
-
+import math
 import random
+import warnings
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, Union
 
@@ -2648,3 +2649,36 @@ def outputs(self) -> Dict[str, Dict[int, str]]:
 
 class PatchTSMixerOnnxConfig(PatchTSTOnnxConfig):
     pass
+
+
+class RTDetrOnnxConfig(ViTOnnxConfig):
+    # Export the operator 'aten::grid_sampler' to ONNX fails under opset 16.
+    # Support for this operator was added in version 16.
+    DEFAULT_ONNX_OPSET = 16
+    ATOL_FOR_VALIDATION = 1e-5
+
+    @property
+    def inputs(self) -> Dict[str, Dict[int, str]]:
+        return {
+            "pixel_values": {0: "batch_size", 2: "height", 3: "width"},
+        }
+
+    def _create_dummy_input_generator_classes(self, **kwargs) -> List["DummyInputGenerator"]:
+        min_image_size = int(math.ceil(self._config.num_queries / 32) * 32)
+        if kwargs["height"] < min_image_size:
+            warnings.warn(
+                f"Exporting model with image `height={kwargs['height']}` which is less than "
+                f"minimal {min_image_size}, setting `height` to {min_image_size}."
+            )
+            kwargs["height"] = min_image_size
+        if kwargs["width"] < min_image_size:
+            warnings.warn(
+                f"Exporting model with image `width={kwargs['width']}` which is less than "
+                f"minimal {min_image_size}, setting `width` to {min_image_size}."
+            )
+            kwargs["width"] = min_image_size
+        return super()._create_dummy_input_generator_classes(**kwargs)
+
+
+class RTDetrV2OnnxConfig(RTDetrOnnxConfig):
+    pass
diff --git a/optimum/exporters/tasks.py b/optimum/exporters/tasks.py
@@ -1104,6 +1104,14 @@ class TasksManager:
             onnx="RoFormerOnnxConfig",
             tflite="RoFormerTFLiteConfig",
         ),
+        "rt-detr": supported_tasks_mapping(
+            "object-detection",
+            onnx="RTDetrOnnxConfig",
+        ),
+        "rt-detr-v2": supported_tasks_mapping(
+            "object-detection",
+            onnx="RTDetrV2OnnxConfig",
+        ),
         "sam": supported_tasks_mapping(
             "feature-extraction",
             onnx="SamOnnxConfig",

diff --git a/tests/exporters/exporters_utils.py b/tests/exporters/exporters_utils.py
@@ -159,6 +159,8 @@
     "resnet": "hf-internal-testing/tiny-random-resnet",
     "roberta": "hf-internal-testing/tiny-random-RobertaModel",
     "roformer": "hf-internal-testing/tiny-random-RoFormerModel",
+    "rt-detr": "PekingU/rtdetr_r18vd",
+    "rt-detr-v2": "PekingU/rtdetr_v2_r18vd",
     "sam": "fxmarty/sam-vit-tiny-random",
     "segformer": "hf-internal-testing/tiny-random-SegformerModel",
     "siglip": "hf-internal-testing/tiny-random-SiglipModel",
@@ -288,6 +290,8 @@
     "resnet": "microsoft/resnet-50",
     "roberta": "roberta-base",
     "roformer": "junnyu/roformer_chinese_base",
+    "rt-detr": "PekingU/rtdetr_r101vd",
+    "rt-detr-v2": "PekingU/rtdetr_v2_r101vd",
     "sam": "facebook/sam-vit-base",
     "segformer": "nvidia/segformer-b0-finetuned-ade-512-512",
     "siglip": "google/siglip-base-patch16-224",