Add EfficientNet Presets (#1933)

* WIP initial conversions

* WIP commit, meant to be able to share current state easily

* WIP before rebase

* WIP rebased on latest master changes

* update verification script to include classifier head

* WIP, 2 variants verified

* added tests, completed 2 variants

* fix efficientnet script examples

* removed copyright heading

* removing unsupported variants for now

* removed explicit cast, corrected tests

* readd explicit cast

* update model uri's

* updated preset uri's

* update projection_activation naming, fix final URI

* update projection_activation variable naming in test

keras_hub/api/layers/__init__.py

@@ -41,6 +41,9 @@
 from keras_hub.src.models.densenet.densenet_image_converter import (
     DenseNetImageConverter,
 )
+from keras_hub.src.models.efficientnet.efficientnet_image_converter import (
+    EfficientNetImageConverter,
+)
 from keras_hub.src.models.mit.mit_image_converter import MiTImageConverter
 from keras_hub.src.models.pali_gemma.pali_gemma_image_converter import (
     PaliGemmaImageConverter,

keras_hub/api/models/__init__.py

@@ -131,6 +131,12 @@
 from keras_hub.src.models.efficientnet.efficientnet_backbone import (
     EfficientNetBackbone,
 )
+from keras_hub.src.models.efficientnet.efficientnet_image_classifier import (
+    EfficientNetImageClassifier,
+)
+from keras_hub.src.models.efficientnet.efficientnet_image_classifier_preprocessor import (
+    EfficientNetImageClassifierPreprocessor,
+)
 from keras_hub.src.models.electra.electra_backbone import ElectraBackbone
 from keras_hub.src.models.electra.electra_tokenizer import ElectraTokenizer
 from keras_hub.src.models.f_net.f_net_backbone import FNetBackbone

keras_hub/src/models/efficientnet/__init__.py

@@ -0,0 +1,9 @@
+from keras_hub.src.models.efficientnet.efficientnet_backbone import (
+    EfficientNetBackbone,
+)
+from keras_hub.src.models.efficientnet.efficientnet_presets import (
+    backbone_presets,
+)
+from keras_hub.src.utils.preset_utils import register_presets
+
+register_presets(backbone_presets, EfficientNetBackbone)

keras_hub/src/models/efficientnet/efficientnet_backbone.py

@@ -104,20 +104,23 @@ def __init__(
         depth_divisor=8,
         min_depth=8,
         input_shape=(None, None, 3),
+        data_format="channels_last",
         activation="swish",
-        include_initial_padding=False,
+        include_stem_padding=True,
         use_depth_divisor_as_min_depth=False,
         cap_round_filter_decrease=False,
-        stem_conv_padding="same",
+        stem_conv_padding="valid",
         batch_norm_momentum=0.9,
+        batch_norm_epsilon=1e-5,
+        projection_activation=None,
         **kwargs,
     ):
         image_input = keras.layers.Input(shape=input_shape)
 
         x = image_input  # Intermediate result.
-        if include_initial_padding:
+        if include_stem_padding:
             x = keras.layers.ZeroPadding2D(
-                padding=self._correct_pad_downsample(x, 3),
+                padding=(1, 1),
                 name="stem_conv_pad",
             )(x)
 
@@ -136,13 +139,15 @@ def __init__(
             kernel_size=3,
             strides=2,
             padding=stem_conv_padding,
+            data_format=data_format,
             use_bias=False,
             kernel_initializer=conv_kernel_initializer(),
             name="stem_conv",
         )(x)
 
         x = keras.layers.BatchNormalization(
             momentum=batch_norm_momentum,
+            epsilon=batch_norm_epsilon,
             name="stem_bn",
         )(x)
         x = keras.layers.Activation(activation, name="stem_activation")(x)
@@ -206,10 +211,13 @@ def __init__(
                         filters_out=output_filters,
                         kernel_size=stackwise_kernel_sizes[i],
                         strides=strides,
+                        data_format=data_format,
                         expand_ratio=stackwise_expansion_ratios[i],
                         se_ratio=squeeze_and_excite_ratio,
                         activation=activation,
+                        projection_activation=projection_activation,
                         dropout=dropout * block_id / blocks,
+                        batch_norm_epsilon=batch_norm_epsilon,
                         name=block_name,
                     )
                 else:
@@ -219,6 +227,7 @@ def __init__(
                         expand_ratio=stackwise_expansion_ratios[i],
                         kernel_size=stackwise_kernel_sizes[i],
                         strides=strides,
+                        data_format=data_format,
                         se_ratio=squeeze_and_excite_ratio,
                         activation=activation,
                         dropout=dropout * block_id / blocks,
@@ -241,15 +250,16 @@ def __init__(
         x = keras.layers.Conv2D(
             filters=top_filters,
             kernel_size=1,
-            padding="same",
             strides=1,
+            padding="same",
+            data_format="channels_last",
             kernel_initializer=conv_kernel_initializer(),
             use_bias=False,
             name="top_conv",
-            data_format="channels_last",
         )(x)
         x = keras.layers.BatchNormalization(
             momentum=batch_norm_momentum,
+            epsilon=batch_norm_epsilon,
             name="top_bn",
         )(x)
         x = keras.layers.Activation(
@@ -268,6 +278,7 @@ def __init__(
         self.dropout = dropout
         self.depth_divisor = depth_divisor
         self.min_depth = min_depth
+        self.data_format = data_format
         self.activation = activation
         self.stackwise_kernel_sizes = stackwise_kernel_sizes
         self.stackwise_num_repeats = stackwise_num_repeats
@@ -280,11 +291,13 @@ def __init__(
         self.stackwise_strides = stackwise_strides
         self.stackwise_block_types = stackwise_block_types
 
-        self.include_initial_padding = include_initial_padding
+        self.include_stem_padding = include_stem_padding
         self.use_depth_divisor_as_min_depth = use_depth_divisor_as_min_depth
         self.cap_round_filter_decrease = cap_round_filter_decrease
         self.stem_conv_padding = stem_conv_padding
         self.batch_norm_momentum = batch_norm_momentum
+        self.batch_norm_epsilon = batch_norm_epsilon
+        self.projection_activation = projection_activation
 
     def get_config(self):
         config = super().get_config()
@@ -305,11 +318,13 @@ def get_config(self):
                 "stackwise_squeeze_and_excite_ratios": self.stackwise_squeeze_and_excite_ratios,
                 "stackwise_strides": self.stackwise_strides,
                 "stackwise_block_types": self.stackwise_block_types,
-                "include_initial_padding": self.include_initial_padding,
+                "include_stem_padding": self.include_stem_padding,
                 "use_depth_divisor_as_min_depth": self.use_depth_divisor_as_min_depth,
                 "cap_round_filter_decrease": self.cap_round_filter_decrease,
                 "stem_conv_padding": self.stem_conv_padding,
                 "batch_norm_momentum": self.batch_norm_momentum,
+                "batch_norm_epsilon": self.batch_norm_epsilon,
+                "projection_activation": self.projection_activation,
             }
         )
         return config
@@ -346,10 +361,13 @@ def _apply_efficientnet_block(
         kernel_size=3,
         strides=1,
         activation="swish",
+        projection_activation=None,
         expand_ratio=1,
         se_ratio=0.0,
         dropout=0.0,
+        batch_norm_epsilon=1e-5,
         name="",
+        data_format="channels_last",
     ):
         """An inverted residual block.
 
@@ -375,12 +393,14 @@ def _apply_efficientnet_block(
                 kernel_size=1,
                 strides=1,
                 padding="same",
+                data_format=data_format,
                 use_bias=False,
                 kernel_initializer=conv_kernel_initializer(),
                 name=name + "expand_conv",
             )(inputs)
             x = keras.layers.BatchNormalization(
                 axis=3,
+                epsilon=batch_norm_epsilon,
                 name=name + "expand_bn",
             )(x)
             x = keras.layers.Activation(
@@ -390,25 +410,23 @@ def _apply_efficientnet_block(
             x = inputs
 
         # Depthwise Convolution
-        if strides == 2:
-            x = keras.layers.ZeroPadding2D(
-                padding=self._correct_pad_downsample(x, kernel_size),
-                name=name + "dwconv_pad",
-            )(x)
-            conv_pad = "valid"
-        else:
-            conv_pad = "same"
-
+        padding_pixels = kernel_size // 2
+        x = keras.layers.ZeroPadding2D(
+            padding=(padding_pixels, padding_pixels),
+            name=name + "dwconv_pad",
+        )(x)
         x = keras.layers.DepthwiseConv2D(
             kernel_size=kernel_size,
             strides=strides,
-            padding=conv_pad,
+            padding="valid",
+            data_format=data_format,
             use_bias=False,
             depthwise_initializer=conv_kernel_initializer(),
             name=name + "dwconv",
         )(x)
         x = keras.layers.BatchNormalization(
             axis=3,
+            epsilon=batch_norm_epsilon,
             name=name + "dwconv_bn",
         )(x)
         x = keras.layers.Activation(
@@ -427,6 +445,7 @@ def _apply_efficientnet_block(
                 filters_se,
                 1,
                 padding="same",
+                data_format=data_format,
                 activation=activation,
                 kernel_initializer=conv_kernel_initializer(),
                 name=name + "se_reduce",
@@ -435,6 +454,7 @@ def _apply_efficientnet_block(
                 filters,
                 1,
                 padding="same",
+                data_format=data_format,
                 activation="sigmoid",
                 kernel_initializer=conv_kernel_initializer(),
                 name=name + "se_expand",
@@ -453,11 +473,13 @@ def _apply_efficientnet_block(
         )(x)
         x = keras.layers.BatchNormalization(
             axis=3,
+            epsilon=batch_norm_epsilon,
             name=name + "project_bn",
         )(x)
-        x = keras.layers.Activation(
-            activation, name=name + "project_activation"
-        )(x)
+        if projection_activation:
+            x = keras.layers.Activation(
+                projection_activation, name=name + "projection_activation"
+            )(x)
 
         if strides == 1 and filters_in == filters_out:
             if dropout > 0:

keras_hub/src/models/efficientnet/efficientnet_backbone_test.py

@@ -73,7 +73,7 @@ def test_valid_call_original_v1(self):
             "depth_coefficient": 1.0,
             "stackwise_block_types": ["v1"] * 7,
             "min_depth": None,
-            "include_initial_padding": True,
+            "include_stem_padding": True,
             "use_depth_divisor_as_min_depth": True,
             "cap_round_filter_decrease": True,
             "stem_conv_padding": "valid",

keras_hub/src/models/efficientnet/efficientnet_image_classifier.py

@@ -0,0 +1,14 @@
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.efficientnet.efficientnet_backbone import (
+    EfficientNetBackbone,
+)
+from keras_hub.src.models.efficientnet.efficientnet_image_classifier_preprocessor import (
+    EfficientNetImageClassifierPreprocessor,
+)
+from keras_hub.src.models.image_classifier import ImageClassifier
+
+
+@keras_hub_export("keras_hub.models.EfficientNetImageClassifier")
+class EfficientNetImageClassifier(ImageClassifier):
+    backbone_cls = EfficientNetBackbone
+    preprocessor_cls = EfficientNetImageClassifierPreprocessor

keras_hub/src/models/efficientnet/efficientnet_image_classifier_preprocessor.py

@@ -0,0 +1,16 @@
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.efficientnet.efficientnet_backbone import (
+    EfficientNetBackbone,
+)
+from keras_hub.src.models.efficientnet.efficientnet_image_converter import (
+    EfficientNetImageConverter,
+)
+from keras_hub.src.models.image_classifier_preprocessor import (
+    ImageClassifierPreprocessor,
+)
+
+
+@keras_hub_export("keras_hub.models.EfficientNetImageClassifierPreprocessor")
+class EfficientNetImageClassifierPreprocessor(ImageClassifierPreprocessor):
+    backbone_cls = EfficientNetBackbone
+    image_converter_cls = EfficientNetImageConverter

keras_hub/src/models/efficientnet/efficientnet_image_classifier_test.py

@@ -0,0 +1,84 @@
+import pytest
+from keras import ops
+
+from keras_hub.src.models.efficientnet.efficientnet_backbone import (
+    EfficientNetBackbone,
+)
+from keras_hub.src.models.efficientnet.efficientnet_image_classifier import (
+    EfficientNetImageClassifier,
+)
+from keras_hub.src.tests.test_case import TestCase
+
+
+class EfficientNetImageClassifierTest(TestCase):
+    def setUp(self):
+        self.images = ops.ones((2, 16, 16, 3))
+        self.labels = [0, 3]
+        backbone = EfficientNetBackbone(
+            width_coefficient=1.0,
+            depth_coefficient=1.0,
+            stackwise_kernel_sizes=[3, 3, 5, 3, 5, 5, 3],
+            stackwise_num_repeats=[1, 2, 2, 3, 3, 4, 1],
+            stackwise_input_filters=[32, 16, 24, 40, 80, 112, 192],
+            stackwise_output_filters=[16, 24, 40, 80, 112, 192, 320],
+            stackwise_expansion_ratios=[1, 6, 6, 6, 6, 6, 6],
+            stackwise_strides=[1, 2, 2, 2, 1, 2, 1],
+            stackwise_block_types=["v1"] * 7,
+            stackwise_squeeze_and_excite_ratios=[0.25] * 7,
+            min_depth=None,
+            include_stem_padding=True,
+            use_depth_divisor_as_min_depth=True,
+            cap_round_filter_decrease=True,
+            stem_conv_padding="valid",
+            batch_norm_momentum=0.9,
+            batch_norm_epsilon=1e-5,
+            dropout=0,
+            projection_activation=None,
+        )
+        self.init_kwargs = {
+            "backbone": backbone,
+            "num_classes": 1000,
+        }
+        self.train_data = (self.images, self.labels)
+
+    def test_classifier_basics(self):
+        pytest.skip(
+            reason="TODO: enable after preprocessor flow is figured out"
+        )
+        self.run_task_test(
+            cls=EfficientNetImageClassifier,
+            init_kwargs=self.init_kwargs,
+            train_data=self.train_data,
+            expected_output_shape=(2, 2),
+        )
+
+    @pytest.mark.large
+    def test_smallest_preset(self):
+        # Test that our forward pass is stable!
+        image_batch = self.load_test_image()[None, ...] / 255.0
+        self.run_preset_test(
+            cls=EfficientNetImageClassifier,
+            preset="efficientnet_b0_ra_imagenet",
+            input_data=image_batch,
+            expected_output_shape=(1, 1000),
+            expected_labels=[85],
+        )
+
+    @pytest.mark.large
+    def test_saved_model(self):
+        self.run_model_saving_test(
+            cls=EfficientNetImageClassifier,
+            init_kwargs=self.init_kwargs,
+            input_data=self.images,
+        )
+
+    @pytest.mark.extra_large
+    def test_all_presets(self):
+        for preset in EfficientNetImageClassifier.presets:
+            self.run_preset_test(
+                cls=EfficientNetImageClassifier,
+                preset=preset,
+                init_kwargs={"num_classes": 2},
+                input_data=self.images,
+                expected_output_shape=(2, 2),
+            )

keras_hub/src/models/efficientnet/efficientnet_image_converter.py

@@ -0,0 +1,10 @@
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.layers.preprocessing.image_converter import ImageConverter
+from keras_hub.src.models.efficientnet.efficientnet_backbone import (
+    EfficientNetBackbone,
+)
+
+
+@keras_hub_export("keras_hub.layers.EfficientNetImageConverter")
+class EfficientNetImageConverter(ImageConverter):
+    backbone_cls = EfficientNetBackbone