add efficientnet onnx config and tiny model tests

Author: Thas-Tayapongsak

optimum/exporters/onnx/model_configs.py

@@ -1061,6 +1061,16 @@ def patch_model_for_export(
         return MgpstrModelPatcher(self, model, model_kwargs=model_kwargs)
 
 
+class EfficientNetOnnxConfig(ViTOnnxConfig):
+    @property
+    def outputs(self) -> Dict[str, Dict[int, str]]:
+        common_outputs = super().outputs 
+
+        if self.task == "image-classification":
+            common_outputs["logits"] = {0: "batch_size", 1: "num_classes"}
+        
+        return common_outputs
+
 class SentenceTransformersTransformerOnnxConfig(TextEncoderOnnxConfig):
     NORMALIZED_CONFIG_CLASS = NormalizedTextConfig
     DEFAULT_ONNX_OPSET = 14  # Some bottleneck transformers models require a specific ONNX opset to be successfully exported. We put a rather high opset here for the export to work for all architectures.

optimum/exporters/tasks.py

@@ -639,6 +639,11 @@ class TasksManager:
             "semantic-segmentation",
             onnx="DptOnnxConfig",
         ),
+        "efficientnet": supported_tasks_mapping(
+            "feature-extraction",
+            "image-classification",
+            onnx="EfficientNetOnnxConfig",
+        ),
         "electra": supported_tasks_mapping(
             "feature-extraction",
             "fill-mask",

tests/exporters/exporters_utils.py

@@ -77,6 +77,10 @@
     "detr": "hf-internal-testing/tiny-random-DetrModel",
     "distilbert": "hf-internal-testing/tiny-random-DistilBertModel",
     "dpt": "hf-internal-testing/tiny-random-DPTModel",
+    "efficientnet": {
+        "hf-internal-testing/tiny-random-EfficientNetModel": ["feature-extraction"],
+        "hf-internal-testing/tiny-random-EfficientNetForImageClassification": ["image-classification"],
+    },
     "electra": "hf-internal-testing/tiny-random-ElectraModel",
     "encoder-decoder": {
         "hf-internal-testing/tiny-random-EncoderDecoderModel-bert-bert": [