From 9e8f8db9dc8372151eb510c7ff0cbd704a8d350f Mon Sep 17 00:00:00 2001 From: Joshua Lochner Date: Sat, 11 Nov 2023 01:19:42 +0000 Subject: [PATCH 1/3] Add depth-estimation w/ dpt --- optimum/exporters/onnx/base.py | 3 +++ optimum/exporters/onnx/model_configs.py | 4 ++++ optimum/exporters/tasks.py | 6 ++++++ tests/exporters/exporters_utils.py | 1 + 4 files changed, 14 insertions(+) diff --git a/optimum/exporters/onnx/base.py b/optimum/exporters/onnx/base.py index 2958d3d920..ef541206ad 100644 --- a/optimum/exporters/onnx/base.py +++ b/optimum/exporters/onnx/base.py @@ -146,6 +146,9 @@ class OnnxConfig(ExportConfig, ABC): "audio-frame-classification": OrderedDict({"logits": {0: "batch_size", 1: "sequence_length"}}), "automatic-speech-recognition": OrderedDict({"logits": {0: "batch_size", 1: "sequence_length"}}), "audio-xvector": OrderedDict({"logits": {0: "batch_size"}, "embeddings": {0: "batch_size"}}), + "depth-estimation": OrderedDict( + {"predicted_depth": {0: "batch_size", 1: "num_channels", 2: "height", 3: "width"}} + ), "document-question-answering": OrderedDict({"logits": {0: "batch_size", 1: "sequence_length"}}), "feature-extraction": OrderedDict({"last_hidden_state": {0: "batch_size", 1: "sequence_length"}}), "fill-mask": OrderedDict({"logits": {0: "batch_size", 1: "sequence_length"}}), diff --git a/optimum/exporters/onnx/model_configs.py b/optimum/exporters/onnx/model_configs.py index fb4d190a2c..19d047d6c5 100644 --- a/optimum/exporters/onnx/model_configs.py +++ b/optimum/exporters/onnx/model_configs.py @@ -743,6 +743,10 @@ class Swin2srOnnxConfig(SwinOnnxConfig): pass +class DptOnnxConfig(ViTOnnxConfig): + pass + + class PoolFormerOnnxConfig(ViTOnnxConfig): NORMALIZED_CONFIG_CLASS = NormalizedVisionConfig ATOL_FOR_VALIDATION = 2e-3 diff --git a/optimum/exporters/tasks.py b/optimum/exporters/tasks.py index 2208a6fe26..ec1a02eb78 100644 --- a/optimum/exporters/tasks.py +++ b/optimum/exporters/tasks.py @@ -164,6 +164,7 @@ class TasksManager: "audio-xvector": "AutoModelForAudioXVector", "automatic-speech-recognition": ("AutoModelForSpeechSeq2Seq", "AutoModelForCTC"), "conversational": ("AutoModelForCausalLM", "AutoModelForSeq2SeqLM"), + "depth-estimation": "AutoModelForDepthEstimation", "feature-extraction": "AutoModel", "fill-mask": "AutoModelForMaskedLM", "image-classification": "AutoModelForImageClassification", @@ -497,6 +498,11 @@ class TasksManager: "feature-extraction", onnx="DonutSwinOnnxConfig", ), + "dpt": supported_tasks_mapping( + "feature-extraction", + "depth-estimation", + onnx="DptOnnxConfig", + ), "electra": supported_tasks_mapping( "feature-extraction", "fill-mask", diff --git a/tests/exporters/exporters_utils.py b/tests/exporters/exporters_utils.py index f17184a1b7..efa34b8599 100644 --- a/tests/exporters/exporters_utils.py +++ b/tests/exporters/exporters_utils.py @@ -65,6 +65,7 @@ "donut-swin": "hf-internal-testing/tiny-random-DonutSwinModel", "detr": "hf-internal-testing/tiny-random-DetrModel", # hf-internal-testing/tiny-random-detr is larger "distilbert": "hf-internal-testing/tiny-random-DistilBertModel", + "dpt": "hf-internal-testing/tiny-random-DPTModel", "electra": "hf-internal-testing/tiny-random-ElectraModel", "encoder-decoder": { "hf-internal-testing/tiny-random-EncoderDecoderModel-bert-bert": [ From 0cbb448b026e61424660a100f36b54ef9f708d25 Mon Sep 17 00:00:00 2001 From: Joshua Lochner Date: Sat, 11 Nov 2023 13:30:31 +0000 Subject: [PATCH 2/3] Fix depth-estimation outputs --- optimum/exporters/onnx/base.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/optimum/exporters/onnx/base.py b/optimum/exporters/onnx/base.py index ef541206ad..8bf93feb29 100644 --- a/optimum/exporters/onnx/base.py +++ b/optimum/exporters/onnx/base.py @@ -146,9 +146,7 @@ class OnnxConfig(ExportConfig, ABC): "audio-frame-classification": OrderedDict({"logits": {0: "batch_size", 1: "sequence_length"}}), "automatic-speech-recognition": OrderedDict({"logits": {0: "batch_size", 1: "sequence_length"}}), "audio-xvector": OrderedDict({"logits": {0: "batch_size"}, "embeddings": {0: "batch_size"}}), - "depth-estimation": OrderedDict( - {"predicted_depth": {0: "batch_size", 1: "num_channels", 2: "height", 3: "width"}} - ), + "depth-estimation": OrderedDict({"predicted_depth": {0: "batch_size", 1: "height", 2: "width"}}), "document-question-answering": OrderedDict({"logits": {0: "batch_size", 1: "sequence_length"}}), "feature-extraction": OrderedDict({"last_hidden_state": {0: "batch_size", 1: "sequence_length"}}), "fill-mask": OrderedDict({"logits": {0: "batch_size", 1: "sequence_length"}}), From a6cc493aef0928f923af699587e2ea23d787a6dc Mon Sep 17 00:00:00 2001 From: Joshua Lochner Date: Sat, 11 Nov 2023 14:00:31 +0000 Subject: [PATCH 3/3] Add GLPN onnx export --- optimum/exporters/onnx/model_configs.py | 4 ++++ optimum/exporters/tasks.py | 5 +++++ tests/exporters/exporters_utils.py | 1 + 3 files changed, 10 insertions(+) diff --git a/optimum/exporters/onnx/model_configs.py b/optimum/exporters/onnx/model_configs.py index 19d047d6c5..f4d50ad58d 100644 --- a/optimum/exporters/onnx/model_configs.py +++ b/optimum/exporters/onnx/model_configs.py @@ -747,6 +747,10 @@ class DptOnnxConfig(ViTOnnxConfig): pass +class GlpnOnnxConfig(ViTOnnxConfig): + pass + + class PoolFormerOnnxConfig(ViTOnnxConfig): NORMALIZED_CONFIG_CLASS = NormalizedVisionConfig ATOL_FOR_VALIDATION = 2e-3 diff --git a/optimum/exporters/tasks.py b/optimum/exporters/tasks.py index ec1a02eb78..7545c72d6c 100644 --- a/optimum/exporters/tasks.py +++ b/optimum/exporters/tasks.py @@ -539,6 +539,11 @@ class TasksManager: onnx="FlaubertOnnxConfig", tflite="FlaubertTFLiteConfig", ), + "glpn": supported_tasks_mapping( + "feature-extraction", + "depth-estimation", + onnx="GlpnOnnxConfig", + ), "gpt2": supported_tasks_mapping( "feature-extraction", "feature-extraction-with-past", diff --git a/tests/exporters/exporters_utils.py b/tests/exporters/exporters_utils.py index efa34b8599..6e43b65e34 100644 --- a/tests/exporters/exporters_utils.py +++ b/tests/exporters/exporters_utils.py @@ -85,6 +85,7 @@ "fxmarty/tiny-testing-falcon-alibi": ["text-generation", "text-generation-with-past"], }, "flaubert": "hf-internal-testing/tiny-random-flaubert", + "glpn": "hf-internal-testing/tiny-random-GLPNModel", "gpt2": "hf-internal-testing/tiny-random-gpt2", "gpt-bigcode": "hf-internal-testing/tiny-random-GPTBigCodeModel", "gpt-neo": "hf-internal-testing/tiny-random-GPTNeoModel",