Skip to content

Commit 8748081

Browse files
authored
Add ONNX export for owlv2 (#1689)
* Add ONNX export for owlv2 * Remove opset=14 requirement
1 parent 64f795d commit 8748081

File tree

4 files changed

+12
-0
lines changed

4 files changed

+12
-0
lines changed

optimum/exporters/onnx/model_configs.py

+4
Original file line numberDiff line numberDiff line change
@@ -1092,6 +1092,10 @@ def outputs(self) -> Dict[str, Dict[int, str]]:
10921092
return outputs
10931093

10941094

1095+
class OwlV2OnnxConfig(OwlViTOnnxConfig):
1096+
pass
1097+
1098+
10951099
class LayoutLMOnnxConfig(TextAndVisionOnnxConfig):
10961100
NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args(
10971101
allow_new=True,

optimum/exporters/tasks.py

+5
Original file line numberDiff line numberDiff line change
@@ -832,6 +832,11 @@ class TasksManager:
832832
"token-classification",
833833
onnx="NystromformerOnnxConfig",
834834
),
835+
"owlv2": supported_tasks_mapping(
836+
"feature-extraction",
837+
"zero-shot-object-detection",
838+
onnx="OwlV2OnnxConfig",
839+
),
835840
"owlvit": supported_tasks_mapping(
836841
"feature-extraction",
837842
"zero-shot-object-detection",

optimum/utils/normalized_config.py

+1
Original file line numberDiff line numberDiff line change
@@ -199,6 +199,7 @@ class NormalizedConfigManager:
199199
'levit',
200200
'mobilebert',
201201
'mobilevit',
202+
'owlv2',
202203
'owlvit',
203204
'perceiver',
204205
'roformer',

tests/exporters/exporters_utils.py

+2
Original file line numberDiff line numberDiff line change
@@ -122,6 +122,7 @@
122122
"mt5": "lewtun/tiny-random-mt5",
123123
"nystromformer": "hf-internal-testing/tiny-random-NystromformerModel",
124124
"opt": "hf-internal-testing/tiny-random-OPTModel",
125+
"owlv2": "hf-internal-testing/tiny-random-Owlv2Model",
125126
"owlvit": "hf-tiny-model-private/tiny-random-OwlViTModel",
126127
"pegasus": "hf-internal-testing/tiny-random-PegasusModel",
127128
"perceiver": {
@@ -243,6 +244,7 @@
243244
"mpt": "mosaicml/mpt-7b",
244245
"mt5": "lewtun/tiny-random-mt5", # Not using google/mt5-small because it takes too much time for testing.
245246
"nystromformer": "hf-internal-testing/tiny-random-NystromformerModel",
247+
"owlv2": "google/owlv2-base-patch16",
246248
"owlvit": "google/owlvit-base-patch32",
247249
"perceiver": "hf-internal-testing/tiny-random-PerceiverModel", # Not using deepmind/language-perceiver because it takes too much time for testing.
248250
# "rembert": "google/rembert",

0 commit comments

Comments
 (0)