add closest variants

pkgoogle · Oct 31, 2024 · 0271f9a · 0271f9a
1 parent 7e39d97
commit 0271f9a
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 0 deletions.
diff --git a/keras_hub/src/utils/timm/convert_efficientnet.py b/keras_hub/src/utils/timm/convert_efficientnet.py
@@ -18,6 +18,22 @@
         "width_coefficient": 1.0,
         "depth_coefficient": 1.1,
     },
+    "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://kerashub/efficientnet/keras/efficientnet_b0_ra_imagenet
+python tools/checkpoint_conversion/convert_efficientnet_checkpoints.py \
+    --preset efficientnet_b0_ra4_e3600_r224_imagenet --upload_uri kaggle://kerashub/efficientnet/keras/efficientnet_b0_ra4_e3600_r224_imagenet
 python tools/checkpoint_conversion/convert_efficientnet_checkpoints.py \
     --preset efficientnet_b1_ft_imagenet --upload_uri kaggle://kerashub/efficientnet/keras/efficientnet_b1_ft_imagenet
+python tools/checkpoint_conversion/convert_efficientnet_checkpoints.py \
+    --preset efficientnet_b1_ra4_e3600_r240_imagenet --upload_uri kaggle://kerashub/efficientnet/keras/efficientnet_b1_ra4_e3600_r240_imagenet
+python tools/checkpoint_conversion/convert_efficientnet_checkpoints.py \
+    --preset efficientnet_b2_ra_imagenet --upload_uri kaggle://kerashub/efficientnet/keras/efficientnet_b2_ra_imagenet
+python tools/checkpoint_conversion/convert_efficientnet_checkpoints.py \
+    --preset efficientnet_b3_ra2_imagenet --upload_uri kaggle://kerashub/efficientnet/keras/efficientnet_b3_ra2_imagenet
+python tools/checkpoint_conversion/convert_efficientnet_checkpoints.py \
+    --preset efficientnet_b4_ra2_imagenet --upload_uri kaggle://kerashub/efficientnet/keras/efficientnet_b4_ra2_imagenet
+python tools/checkpoint_conversion/convert_efficientnet_checkpoints.py \
+    --preset efficientnet_b5_sw_imagenet --upload_uri kaggle://kerashub/efficientnet/keras/efficientnet_b5_sw_imagenet
+python tools/checkpoint_conversion/convert_efficientnet_checkpoints.py \
+    --preset efficientnet_b5_sw_ft_imagenet --upload_uri kaggle://kerashub/efficientnet/keras/efficientnet_b5_sw_ft_imagenet
 """
 
 import os
@@ -22,7 +36,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",
 }
 FLAGS = flags.FLAGS