From 4025500f21e2bdb69f631ec270387c7bfdc604ad Mon Sep 17 00:00:00 2001
From: pkgoogle <132095473+pkgoogle@users.noreply.github.com>
Date: Wed, 20 Nov 2024 20:35:43 +0000
Subject: [PATCH] EfficientNet: Add lite0 variant (#1968)

* add lite0 variant

* format pass

* added to preset file, fix uri

* fix parameter regression

* format fix

* fix merge regression
---
 .../efficientnet/efficientnet_presets.py      | 13 ++++++
 .../src/utils/timm/convert_efficientnet.py    | 42 ++++++++++++-------
 .../convert_efficientnet_checkpoints.py       |  3 ++
 3 files changed, 42 insertions(+), 16 deletions(-)

diff --git a/keras_hub/src/models/efficientnet/efficientnet_presets.py b/keras_hub/src/models/efficientnet/efficientnet_presets.py
index 0807f7e163..9e82a7ca5d 100644
--- a/keras_hub/src/models/efficientnet/efficientnet_presets.py
+++ b/keras_hub/src/models/efficientnet/efficientnet_presets.py
@@ -163,4 +163,17 @@
         },
         "kaggle_handle": "kaggle://keras/efficientnet/keras/efficientnet_es_ra_imagenet/1",
     },
+    "efficientnet_lite0_ra_imagenet": {
+        "metadata": {
+            "description": (
+                "EfficientNet-Lite model fine-trained on the ImageNet 1k dataset "
+                "with RandAugment recipe."
+            ),
+            "params": 4652008,
+            "official_name": "EfficientNet",
+            "path": "efficientnet",
+            "model_card": "https://arxiv.org/abs/1905.11946",
+        },
+        "kaggle_handle": "kaggle://keras/efficientnet/keras/efficientnet_lite0_ra_imagenet",
+    },
 }
diff --git a/keras_hub/src/utils/timm/convert_efficientnet.py b/keras_hub/src/utils/timm/convert_efficientnet.py
index 5bfa8dbec6..0b2c1edbae 100644
--- a/keras_hub/src/utils/timm/convert_efficientnet.py
+++ b/keras_hub/src/utils/timm/convert_efficientnet.py
@@ -20,6 +20,32 @@
         "depth_coefficient": 1.1,
         "stackwise_squeeze_and_excite_ratios": [0.25] * 7,
     },
+    "b2": {
+        "width_coefficient": 1.1,
+        "depth_coefficient": 1.2,
+        "stackwise_squeeze_and_excite_ratios": [0.25] * 7,
+    },
+    "b3": {
+        "width_coefficient": 1.2,
+        "depth_coefficient": 1.4,
+        "stackwise_squeeze_and_excite_ratios": [0.25] * 7,
+    },
+    "b4": {
+        "width_coefficient": 1.4,
+        "depth_coefficient": 1.8,
+        "stackwise_squeeze_and_excite_ratios": [0.25] * 7,
+    },
+    "b5": {
+        "width_coefficient": 1.6,
+        "depth_coefficient": 2.2,
+        "stackwise_squeeze_and_excite_ratios": [0.25] * 7,
+    },
+    "lite0": {
+        "width_coefficient": 1.0,
+        "depth_coefficient": 1.0,
+        "stackwise_squeeze_and_excite_ratios": [0] * 7,
+        "activation": "relu6",
+    },
     "el": {
         "width_coefficient": 1.2,
         "depth_coefficient": 1.4,
@@ -65,22 +91,6 @@
         "stackwise_nores_option": [True] + [False] * 5,
         "activation": "relu",
     },
-    "b2": {
-        "width_coefficient": 1.1,
-        "depth_coefficient": 1.2,
-    },
-    "b3": {
-        "width_coefficient": 1.2,
-        "depth_coefficient": 1.4,
-    },
-    "b4": {
-        "width_coefficient": 1.4,
-        "depth_coefficient": 1.8,
-    },
-    "b5": {
-        "width_coefficient": 1.6,
-        "depth_coefficient": 2.2,
-    },
 }
 
 
diff --git a/tools/checkpoint_conversion/convert_efficientnet_checkpoints.py b/tools/checkpoint_conversion/convert_efficientnet_checkpoints.py
index aac681a7a7..53f3c90297 100644
--- a/tools/checkpoint_conversion/convert_efficientnet_checkpoints.py
+++ b/tools/checkpoint_conversion/convert_efficientnet_checkpoints.py
@@ -19,6 +19,8 @@
     --preset efficientnet_b5_sw_imagenet --upload_uri kaggle://keras/efficientnet/keras/efficientnet_b5_sw_imagenet
 python tools/checkpoint_conversion/convert_efficientnet_checkpoints.py \
     --preset efficientnet_b5_sw_ft_imagenet --upload_uri kaggle://keras/efficientnet/keras/efficientnet_b5_sw_ft_imagenet
+python tools/checkpoint_conversion/convert_efficientnet_checkpoints.py \
+    --preset efficientnet_lite0_ra_imagenet --upload_uri kaggle://keras/efficientnet/keras/efficientnet_lite0_ra_imagenet
 python tools/checkpoint_conversion/convert_efficientnet_checkpoints.py \
     --preset efficientnet_el_ra_imagenet --upload_uri kaggle://keras/efficientnet/keras/efficientnet_el_ra_imagenet
 python tools/checkpoint_conversion/convert_efficientnet_checkpoints.py \
@@ -50,6 +52,7 @@
     "efficientnet_b4_ra2_imagenet": "timm/efficientnet_b4.ra2_in1k",
     "efficientnet_b5_sw_imagenet": "timm/efficientnet_b5.sw_in12k",
     "efficientnet_b5_sw_ft_imagenet": "timm/efficientnet_b5.sw_in12k_ft_in1k",
+    "efficientnet_lite0_ra_imagenet": "timm/efficientnet_lite0.ra_in1k",
     "efficientnet_el_ra_imagenet": "timm/efficientnet_el.ra_in1k",
     "efficientnet_em_ra2_imagenet": "timm/efficientnet_em.ra2_in1k",
     "efficientnet_es_ra_imagenet": "timm/efficientnet_es.ra_in1k",