ORTOptimizer for the model type Segformer

optimum/onnxruntime/utils.py

@@ -20,12 +20,11 @@
 from typing import Callable, Dict, List, Optional, Tuple, Union
 
 import numpy as np
+import onnxruntime as ort
 import torch
 from packaging import version
 from transformers.utils import logging
 
-import onnxruntime as ort
-
 from ..exporters.onnx import OnnxConfig, OnnxConfigWithLoss
 from ..utils.import_utils import _is_package_available
 
@@ -128,6 +127,7 @@ class ORTConfigManager:
         "nystromformer": "bert",
         "pegasus": "bert",
         "roberta": "bert",
+        "segformer": "vit",
         "t5": "bert",
         "vit": "vit",
         "whisper": "bart",

optimum/utils/normalized_config.py

@@ -102,6 +102,30 @@ class NormalizedVisionConfig(NormalizedConfig):
     INPUT_SIZE = "input_size"
 
 
+class NormalizedSegformerConfig(NormalizedVisionConfig):
+    NUM_ATTENTION_HEADS = "num_attention_heads"
+    HIDDEN_SIZE = "decoder_hidden_size"
+
+    def __getattr__(self, attr_name):
+        if attr_name == "hidden_size":
+            attr_value = getattr(self.config, self.HIDDEN_SIZE, None)
+            if attr_value is None:
+                raise AttributeError(f"Attribute {self.HIDDEN_SIZE} not found in config")
+            if isinstance(attr_value, list):
+                return max(attr_value)
+            return attr_value
+
+        elif attr_name == "num_attention_heads":
+            attr_value = getattr(self.config, self.NUM_ATTENTION_HEADS, None)
+            if attr_value is None:
+                raise AttributeError(f"Attribute {self.NUM_ATTENTION_HEADS} not found in config")
+            if isinstance(attr_value, list):
+                return max(attr_value)
+            return attr_value
+
+        return super().__getattr__(attr_name)
+
+
 class NormalizedTextAndVisionConfig(NormalizedTextConfig, NormalizedVisionConfig):
     TEXT_CONFIG = None
     VISION_CONFIG = None
@@ -203,7 +227,6 @@ class NormalizedConfigManager:
         'owlvit',
         'perceiver',
         'roformer',
-        'segformer',
         'squeezebert',
         'table-transformer',
     """
@@ -256,6 +279,7 @@ class NormalizedConfigManager:
         "regnet": NormalizedVisionConfig,
         "resnet": NormalizedVisionConfig,
         "roberta": NormalizedTextConfig,
+        "segformer": NormalizedSegformerConfig,
         "speech-to-text": SpeechToTextLikeNormalizedTextConfig,
         "splinter": NormalizedTextConfig,
         "t5": T5LikeNormalizedTextConfig,

tests/onnxruntime/test_optimization.py

@@ -24,12 +24,6 @@
 import onnx
 import pytest
 import torch
-from parameterized import parameterized
-from transformers import AutoTokenizer
-from transformers.onnx.utils import get_preprocessor
-from transformers.testing_utils import require_torch_gpu
-from utils_onnxruntime_tests import MODEL_NAMES
-
 from optimum.exporters import TasksManager
 from optimum.exporters.onnx import MODEL_TYPES_REQUIRING_POSITION_IDS
 from optimum.onnxruntime import (
@@ -43,6 +37,11 @@
 from optimum.onnxruntime.modeling_decoder import ORTModelForCausalLM
 from optimum.onnxruntime.modeling_seq2seq import ORTModelForSeq2SeqLM, ORTModelForSpeechSeq2Seq
 from optimum.utils.testing_utils import grid_parameters
+from parameterized import parameterized
+from transformers import AutoTokenizer
+from transformers.onnx.utils import get_preprocessor
+from transformers.testing_utils import require_torch_gpu
+from utils_onnxruntime_tests import MODEL_NAMES
 
 
 class ORTOptimizerTestMixin(unittest.TestCase):
@@ -171,6 +170,7 @@ def test_compare_original_seq2seq_model_with_optimized_model(self, model_cls, mo
 
     # Contribution note: Please add test models in alphabetical order. Find test models here: https://huggingface.co/hf-internal-testing.
     SUPPORTED_IMAGE_ARCHITECTURES_WITH_MODEL_ID = (
+        (ORTModelForImageClassification, "hf-internal-testing/tiny-random-segformer"),
         (ORTModelForImageClassification, "hf-internal-testing/tiny-random-vit"),
     )