Add closest EfficientNet variants (#1967)

* add closest variants * added to presets file * format pass * fix quote literal and formatting * fix uris * add version to kaggle handles * fix preset config * fix description error
keras-team · Nov 14, 2024 · 123ef82 · 123ef82
1 parent 0756fb4
commit 123ef82
Show file tree

Hide file tree

Showing 3 changed files with 135 additions and 0 deletions.
diff --git a/keras_hub/src/models/efficientnet/efficientnet_presets.py b/keras_hub/src/models/efficientnet/efficientnet_presets.py
@@ -14,6 +14,21 @@
         },
         "kaggle_handle": "kaggle://keras/efficientnet/keras/efficientnet_b0_ra_imagenet/1",
     },
+    "efficientnet_b0_ra4_e3600_r224_imagenet": {
+        "metadata": {
+            "description": (
+                "EfficientNet B0 model pre-trained on the ImageNet 1k dataset by"
+                " Ross Wightman. Trained with timm scripts using hyper-parameters"
+                " inspired by the MobileNet-V4 small, mixed with go-to hparams "
+                'from timm and "ResNet Strikes Back".'
+            ),
+            "params": 5288548,
+            "official_name": "EfficientNet",
+            "path": "efficientnet",
+            "model_card": "https://arxiv.org/abs/1905.11946",
+        },
+        "kaggle_handle": "kaggle://keras/efficientnet/keras/efficientnet_b0_ra4_e3600_r224_imagenet/1",
+    },
     "efficientnet_b1_ft_imagenet": {
         "metadata": {
             "description": (
@@ -26,6 +41,89 @@
         },
         "kaggle_handle": "kaggle://keras/efficientnet/keras/efficientnet_b1_ft_imagenet/1",
     },
+    "efficientnet_b1_ra4_e3600_r240_imagenet": {
+        "metadata": {
+            "description": (
+                "EfficientNet B1 model pre-trained on the ImageNet 1k dataset by"
+                " Ross Wightman. Trained with timm scripts using hyper-parameters"
+                " inspired by the MobileNet-V4 small, mixed with go-to hparams "
+                'from timm and "ResNet Strikes Back".'
+            ),
+            "params": 7794184,
+            "official_name": "EfficientNet",
+            "path": "efficientnet",
+            "model_card": "https://arxiv.org/abs/1905.11946",
+        },
+        "kaggle_handle": "kaggle://keras/efficientnet/keras/efficientnet_b1_ra4_e3600_r240_imagenet/1",
+    },
+    "efficientnet_b2_ra_imagenet": {
+        "metadata": {
+            "description": (
+                "EfficientNet B2 model pre-trained on the ImageNet 1k dataset "
+                "with RandAugment recipe."
+            ),
+            "params": 9109994,
+            "official_name": "EfficientNet",
+            "path": "efficientnet",
+            "model_card": "https://arxiv.org/abs/1905.11946",
+        },
+        "kaggle_handle": "kaggle://keras/efficientnet/keras/efficientnet_b2_ra_imagenet/1",
+    },
+    "efficientnet_b3_ra2_imagenet": {
+        "metadata": {
+            "description": (
+                "EfficientNet B3 model pre-trained on the ImageNet 1k dataset "
+                "with RandAugment2 recipe."
+            ),
+            "params": 12233232,
+            "official_name": "EfficientNet",
+            "path": "efficientnet",
+            "model_card": "https://arxiv.org/abs/1905.11946",
+        },
+        "kaggle_handle": "kaggle://keras/efficientnet/keras/efficientnet_b3_ra2_imagenet/1",
+    },
+    "efficientnet_b4_ra2_imagenet": {
+        "metadata": {
+            "description": (
+                "EfficientNet B4 model pre-trained on the ImageNet 1k dataset "
+                "with RandAugment2 recipe."
+            ),
+            "params": 19341616,
+            "official_name": "EfficientNet",
+            "path": "efficientnet",
+            "model_card": "https://arxiv.org/abs/1905.11946",
+        },
+        "kaggle_handle": "kaggle://keras/efficientnet/keras/efficientnet_b4_ra2_imagenet/1",
+    },
+    "efficientnet_b5_sw_imagenet": {
+        "metadata": {
+            "description": (
+                "EfficientNet B5 model pre-trained on the ImageNet 12k dataset "
+                "by Ross Wightman. Based on Swin Transformer train / pretrain "
+                "recipe with modifications (related to both DeiT and ConvNeXt recipes)."
+            ),
+            "params": 30389784,
+            "official_name": "EfficientNet",
+            "path": "efficientnet",
+            "model_card": "https://arxiv.org/abs/1905.11946",
+        },
+        "kaggle_handle": "kaggle://keras/efficientnet/keras/efficientnet_b5_sw_imagenet/1",
+    },
+    "efficientnet_b5_sw_ft_imagenet": {
+        "metadata": {
+            "description": (
+                "EfficientNet B5 model pre-trained on the ImageNet 12k dataset "
+                "and fine-tuned on ImageNet-1k by Ross Wightman. Based on Swin "
+                "Transformer train / pretrain recipe with modifications "
+                "(related to both DeiT and ConvNeXt recipes)."
+            ),
+            "params": 30389784,
+            "official_name": "EfficientNet",
+            "path": "efficientnet",
+            "model_card": "https://arxiv.org/abs/1905.11946",
+        },
+        "kaggle_handle": "kaggle://keras/efficientnet/keras/efficientnet_b5_sw_ft_imagenet/1",
+    },
     "efficientnet_el_ra_imagenet": {
         "metadata": {
             "description": (

diff --git a/keras_hub/src/utils/timm/convert_efficientnet.py b/keras_hub/src/utils/timm/convert_efficientnet.py
@@ -65,6 +65,22 @@
         "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
@@ -3,8 +3,22 @@
 
 python tools/checkpoint_conversion/convert_efficientnet_checkpoints.py \
     --preset efficientnet_b0_ra_imagenet --upload_uri kaggle://keras/efficientnet/keras/efficientnet_b0_ra_imagenet
+python tools/checkpoint_conversion/convert_efficientnet_checkpoints.py \
+    --preset efficientnet_b0_ra4_e3600_r224_imagenet --upload_uri kaggle://keras/efficientnet/keras/efficientnet_b0_ra4_e3600_r224_imagenet
 python tools/checkpoint_conversion/convert_efficientnet_checkpoints.py \
     --preset efficientnet_b1_ft_imagenet --upload_uri kaggle://keras/efficientnet/keras/efficientnet_b1_ft_imagenet
+python tools/checkpoint_conversion/convert_efficientnet_checkpoints.py \
+    --preset efficientnet_b1_ra4_e3600_r240_imagenet --upload_uri kaggle://keras/efficientnet/keras/efficientnet_b1_ra4_e3600_r240_imagenet
+python tools/checkpoint_conversion/convert_efficientnet_checkpoints.py \
+    --preset efficientnet_b2_ra_imagenet --upload_uri kaggle://keras/efficientnet/keras/efficientnet_b2_ra_imagenet
+python tools/checkpoint_conversion/convert_efficientnet_checkpoints.py \
+    --preset efficientnet_b3_ra2_imagenet --upload_uri kaggle://keras/efficientnet/keras/efficientnet_b3_ra2_imagenet
+python tools/checkpoint_conversion/convert_efficientnet_checkpoints.py \
+    --preset efficientnet_b4_ra2_imagenet --upload_uri kaggle://keras/efficientnet/keras/efficientnet_b4_ra2_imagenet
+python tools/checkpoint_conversion/convert_efficientnet_checkpoints.py \
+    --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_el_ra_imagenet --upload_uri kaggle://keras/efficientnet/keras/efficientnet_el_ra_imagenet
 python tools/checkpoint_conversion/convert_efficientnet_checkpoints.py \
@@ -28,7 +42,14 @@
 
 PRESET_MAP = {
     "efficientnet_b0_ra_imagenet": "timm/efficientnet_b0.ra_in1k",
+    "efficientnet_b0_ra4_e3600_r224_imagenet": "timm/efficientnet_b0.ra4_e3600_r224_in1k",
     "efficientnet_b1_ft_imagenet": "timm/efficientnet_b1.ft_in1k",
+    "efficientnet_b1_ra4_e3600_r240_imagenet": "timm/efficientnet_b1.ra4_e3600_r240_in1k",
+    "efficientnet_b2_ra_imagenet": "timm/efficientnet_b2.ra_in1k",
+    "efficientnet_b3_ra2_imagenet": "timm/efficientnet_b3.ra2_in1k",
+    "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_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",