Add ONNX export for owlv2 (#1689)

* Add ONNX export for owlv2 * Remove opset=14 requirement
huggingface · Feb 29, 2024 · 8748081 · 8748081
1 parent 64f795d
commit 8748081
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 0 deletions.
diff --git a/optimum/exporters/onnx/model_configs.py b/optimum/exporters/onnx/model_configs.py
@@ -1092,6 +1092,10 @@ def outputs(self) -> Dict[str, Dict[int, str]]:
         return outputs
 
 
+class OwlV2OnnxConfig(OwlViTOnnxConfig):
+    pass
+
+
 class LayoutLMOnnxConfig(TextAndVisionOnnxConfig):
     NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args(
         allow_new=True,

diff --git a/optimum/exporters/tasks.py b/optimum/exporters/tasks.py
@@ -832,6 +832,11 @@ class TasksManager:
             "token-classification",
             onnx="NystromformerOnnxConfig",
         ),
+        "owlv2": supported_tasks_mapping(
+            "feature-extraction",
+            "zero-shot-object-detection",
+            onnx="OwlV2OnnxConfig",
+        ),
         "owlvit": supported_tasks_mapping(
             "feature-extraction",
             "zero-shot-object-detection",

diff --git a/optimum/utils/normalized_config.py b/optimum/utils/normalized_config.py
@@ -199,6 +199,7 @@ class NormalizedConfigManager:
         'levit',
         'mobilebert',
         'mobilevit',
+        'owlv2',
         'owlvit',
         'perceiver',
         'roformer',

diff --git a/tests/exporters/exporters_utils.py b/tests/exporters/exporters_utils.py
@@ -122,6 +122,7 @@
     "mt5": "lewtun/tiny-random-mt5",
     "nystromformer": "hf-internal-testing/tiny-random-NystromformerModel",
     "opt": "hf-internal-testing/tiny-random-OPTModel",
+    "owlv2": "hf-internal-testing/tiny-random-Owlv2Model",
     "owlvit": "hf-tiny-model-private/tiny-random-OwlViTModel",
     "pegasus": "hf-internal-testing/tiny-random-PegasusModel",
     "perceiver": {
@@ -243,6 +244,7 @@
     "mpt": "mosaicml/mpt-7b",
     "mt5": "lewtun/tiny-random-mt5",  # Not using google/mt5-small because it takes too much time for testing.
     "nystromformer": "hf-internal-testing/tiny-random-NystromformerModel",
+    "owlv2": "google/owlv2-base-patch16",
     "owlvit": "google/owlvit-base-patch32",
     "perceiver": "hf-internal-testing/tiny-random-PerceiverModel",  # Not using deepmind/language-perceiver because it takes too much time for testing.
     # "rembert": "google/rembert",