keras-hub-nightly 0.16.1.dev202410200345__py3-none-any.whl → 0.19.0.dev202412070351__py3-none-any.whl

keras_hub/src/models/efficientnet/efficientnet_presets.py

@@ -0,0 +1,192 @@
+"""EfficientNet preset configurations."""
+
+backbone_presets = {
+    "efficientnet_b0_ra_imagenet": {
+        "metadata": {
+            "description": (
+                "EfficientNet B0 model pre-trained on the ImageNet 1k dataset "
+                "with RandAugment recipe."
+            ),
+            "params": 5288548,
+            "path": "efficientnet",
+        },
+        "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,
+            "path": "efficientnet",
+        },
+        "kaggle_handle": "kaggle://keras/efficientnet/keras/efficientnet_b0_ra4_e3600_r224_imagenet/1",
+    },
+    "efficientnet_b1_ft_imagenet": {
+        "metadata": {
+            "description": (
+                "EfficientNet B1 model fine-tuned on the ImageNet 1k dataset."
+            ),
+            "params": 7794184,
+            "path": "efficientnet",
+        },
+        "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,
+            "path": "efficientnet",
+        },
+        "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,
+            "path": "efficientnet",
+        },
+        "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,
+            "path": "efficientnet",
+        },
+        "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,
+            "path": "efficientnet",
+        },
+        "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,
+            "path": "efficientnet",
+        },
+        "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,
+            "path": "efficientnet",
+        },
+        "kaggle_handle": "kaggle://keras/efficientnet/keras/efficientnet_b5_sw_ft_imagenet/1",
+    },
+    "efficientnet_el_ra_imagenet": {
+        "metadata": {
+            "description": (
+                "EfficientNet-EdgeTPU Large model trained on the ImageNet 1k "
+                "dataset with RandAugment recipe."
+            ),
+            "params": 10589712,
+            "path": "efficientnet",
+        },
+        "kaggle_handle": "kaggle://keras/efficientnet/keras/efficientnet_b1_ft_imagenet",
+    },
+    "efficientnet_em_ra2_imagenet": {
+        "metadata": {
+            "description": (
+                "EfficientNet-EdgeTPU Medium model trained on the ImageNet 1k "
+                "dataset with RandAugment2 recipe."
+            ),
+            "params": 6899496,
+            "path": "efficientnet",
+        },
+        "kaggle_handle": "kaggle://keras/efficientnet/keras/efficientnet_b1_ft_imagenet",
+    },
+    "efficientnet_es_ra_imagenet": {
+        "metadata": {
+            "description": (
+                "EfficientNet-EdgeTPU Small model trained on the ImageNet 1k "
+                "dataset with RandAugment recipe."
+            ),
+            "params": 5438392,
+            "path": "efficientnet",
+        },
+        "kaggle_handle": "kaggle://keras/efficientnet/keras/efficientnet_b1_ft_imagenet",
+    },
+    "efficientnet2_rw_m_agc_imagenet": {
+        "metadata": {
+            "description": (
+                "EfficientNet-v2 Medium model trained on the ImageNet 1k "
+                "dataset with adaptive gradient clipping."
+            ),
+            "params": 53236442,
+            "official_name": "EfficientNet",
+            "path": "efficientnet",
+            "model_card": "https://arxiv.org/abs/2104.00298",
+        },
+        "kaggle_handle": "kaggle://keras/efficientnet/keras/efficientnet2_rw_m_agc_imagenet",
+    },
+    "efficientnet2_rw_s_ra2_imagenet": {
+        "metadata": {
+            "description": (
+                "EfficientNet-v2 Small model trained on the ImageNet 1k "
+                "dataset with RandAugment2 recipe."
+            ),
+            "params": 23941296,
+            "official_name": "EfficientNet",
+            "path": "efficientnet",
+            "model_card": "https://arxiv.org/abs/2104.00298",
+        },
+        "kaggle_handle": "kaggle://keras/efficientnet/keras/efficientnet2_rw_s_ra2_imagenet",
+    },
+    "efficientnet2_rw_t_ra2_imagenet": {
+        "metadata": {
+            "description": (
+                "EfficientNet-v2 Tiny model trained on the ImageNet 1k "
+                "dataset with RandAugment2 recipe."
+            ),
+            "params": 13649388,
+            "official_name": "EfficientNet",
+            "path": "efficientnet",
+            "model_card": "https://arxiv.org/abs/2104.00298",
+        },
+        "kaggle_handle": "kaggle://keras/efficientnet/keras/efficientnet2_rw_t_ra2_imagenet",
+    },
+    "efficientnet_lite0_ra_imagenet": {
+        "metadata": {
+            "description": (
+                "EfficientNet-Lite model fine-trained on the ImageNet 1k dataset "
+                "with RandAugment recipe."
+            ),
+            "params": 4652008,
+            "path": "efficientnet",
+        },
+        "kaggle_handle": "kaggle://keras/efficientnet/keras/efficientnet_lite0_ra_imagenet",
+    },
+}

keras_hub/src/models/efficientnet/fusedmbconv.py

@@ -2,15 +2,6 @@ import keras
 
 BN_AXIS = 3
 
-CONV_KERNEL_INITIALIZER = {
-    "class_name": "VarianceScaling",
-    "config": {
-        "scale": 2.0,
-        "mode": "fan_out",
-        "distribution": "truncated_normal",
-    },
-}
-
 
 class FusedMBConvBlock(keras.layers.Layer):
     """Implementation of the FusedMBConv block
@@ -44,13 +35,24 @@ class FusedMBConvBlock(keras.layers.Layer):
             convolutions
         strides: default 1, the strides to apply to the expansion phase
             convolutions
+        data_format: str, channels_last (default) or channels_first, expects
+            tensors to be of shape (N, H, W, C) or (N, C, H, W) respectively
         se_ratio: default 0.0, The filters used in the Squeeze-Excitation phase,
             and are chosen as the maximum between 1 and input_filters*se_ratio
         batch_norm_momentum: default 0.9, the BatchNormalization momentum
+        batch_norm_epsilon: default 1e-3, float, epsilon for batch norm
+            calcualtions. Used in denominator for calculations to prevent divide
+            by 0 errors.
         activation: default "swish", the activation function used between
             convolution operations
+        projection_activation: default None, the activation function to use
+            after the output projection convoultion
         dropout: float, the optional dropout rate to apply before the output
             convolution, defaults to 0.2
+        nores: bool, default False, forces no residual connection if True,
+            otherwise allows it if False.
+        projection_kernel_size: default 1, the kernel_size to apply to the
+            output projection phase convolution
 
     Returns:
         A tensor representing a feature map, passed through the FusedMBConv
@@ -67,10 +69,15 @@ class FusedMBConvBlock(keras.layers.Layer):
         expand_ratio=1,
         kernel_size=3,
         strides=1,
+        data_format="channels_last",
         se_ratio=0.0,
         batch_norm_momentum=0.9,
+        batch_norm_epsilon=1e-3,
         activation="swish",
+        projection_activation=None,
         dropout=0.2,
+        nores=False,
+        projection_kernel_size=1,
         **kwargs
     ):
         super().__init__(**kwargs)
@@ -79,44 +86,50 @@ class FusedMBConvBlock(keras.layers.Layer):
         self.expand_ratio = expand_ratio
         self.kernel_size = kernel_size
         self.strides = strides
+        self.data_format = data_format
         self.se_ratio = se_ratio
         self.batch_norm_momentum = batch_norm_momentum
+        self.batch_norm_epsilon = batch_norm_epsilon
         self.activation = activation
+        self.projection_activation = projection_activation
         self.dropout = dropout
+        self.nores = nores
+        self.projection_kernel_size = projection_kernel_size
         self.filters = self.input_filters * self.expand_ratio
         self.filters_se = max(1, int(input_filters * se_ratio))
 
+        padding_pixels = kernel_size // 2
+        self.conv1_pad = keras.layers.ZeroPadding2D(
+            padding=(padding_pixels, padding_pixels),
+            name=self.name + "expand_conv_pad",
+        )
         self.conv1 = keras.layers.Conv2D(
             filters=self.filters,
             kernel_size=kernel_size,
             strides=strides,
-            kernel_initializer=CONV_KERNEL_INITIALIZER,
-            padding="same",
-            data_format="channels_last",
+            kernel_initializer=self._conv_kernel_initializer(),
+            padding="valid",
+            data_format=data_format,
             use_bias=False,
             name=self.name + "expand_conv",
         )
         self.bn1 = keras.layers.BatchNormalization(
             axis=BN_AXIS,
             momentum=self.batch_norm_momentum,
+            epsilon=self.batch_norm_epsilon,
             name=self.name + "expand_bn",
         )
         self.act = keras.layers.Activation(
             self.activation, name=self.name + "expand_activation"
         )
 
-        self.bn2 = keras.layers.BatchNormalization(
-            axis=BN_AXIS,
-            momentum=self.batch_norm_momentum,
-            name=self.name + "bn",
-        )
-
         self.se_conv1 = keras.layers.Conv2D(
             self.filters_se,
             1,
             padding="same",
+            data_format=data_format,
             activation=self.activation,
-            kernel_initializer=CONV_KERNEL_INITIALIZER,
+            kernel_initializer=self._conv_kernel_initializer(),
             name=self.name + "se_reduce",
         )
 
@@ -124,28 +137,40 @@ class FusedMBConvBlock(keras.layers.Layer):
             self.filters,
             1,
             padding="same",
+            data_format=data_format,
             activation="sigmoid",
-            kernel_initializer=CONV_KERNEL_INITIALIZER,
+            kernel_initializer=self._conv_kernel_initializer(),
             name=self.name + "se_expand",
         )
 
+        padding_pixels = projection_kernel_size // 2
+        self.output_conv_pad = keras.layers.ZeroPadding2D(
+            padding=(padding_pixels, padding_pixels),
+            name=self.name + "project_conv_pad",
+        )
         self.output_conv = keras.layers.Conv2D(
             filters=self.output_filters,
-            kernel_size=1 if expand_ratio != 1 else kernel_size,
+            kernel_size=projection_kernel_size,
             strides=1,
-            kernel_initializer=CONV_KERNEL_INITIALIZER,
-            padding="same",
-            data_format="channels_last",
+            kernel_initializer=self._conv_kernel_initializer(),
+            padding="valid",
+            data_format=data_format,
             use_bias=False,
             name=self.name + "project_conv",
         )
 
-        self.bn3 = keras.layers.BatchNormalization(
+        self.bn2 = keras.layers.BatchNormalization(
             axis=BN_AXIS,
             momentum=self.batch_norm_momentum,
+            epsilon=self.batch_norm_epsilon,
             name=self.name + "project_bn",
         )
 
+        if self.projection_activation:
+            self.projection_act = keras.layers.Activation(
+                self.projection_activation, name=self.name + "projection_act"
+            )
+
         if self.dropout:
             self.dropout_layer = keras.layers.Dropout(
                 self.dropout,
@@ -153,23 +178,33 @@ class FusedMBConvBlock(keras.layers.Layer):
                 name=self.name + "drop",
             )
 
+    def _conv_kernel_initializer(
+        self,
+        scale=2.0,
+        mode="fan_out",
+        distribution="truncated_normal",
+        seed=None,
+    ):
+        return keras.initializers.VarianceScaling(
+            scale=scale, mode=mode, distribution=distribution, seed=seed
+        )
+
     def build(self, input_shape):
         if self.name is None:
             self.name = keras.backend.get_uid("block0")
 
     def call(self, inputs):
         # Expansion phase
-        if self.expand_ratio != 1:
-            x = self.conv1(inputs)
-            x = self.bn1(x)
-            x = self.act(x)
-        else:
-            x = inputs
+        x = self.conv1_pad(inputs)
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.act(x)
 
         # Squeeze and excite
         if 0 < self.se_ratio <= 1:
             se = keras.layers.GlobalAveragePooling2D(
-                name=self.name + "se_squeeze"
+                name=self.name + "se_squeeze",
+                data_format=self.data_format,
             )(x)
             if BN_AXIS == 1:
                 se_shape = (self.filters, 1, 1)
@@ -186,13 +221,18 @@ class FusedMBConvBlock(keras.layers.Layer):
             x = keras.layers.multiply([x, se], name=self.name + "se_excite")
 
         # Output phase:
+        x = self.output_conv_pad(x)
         x = self.output_conv(x)
-        x = self.bn3(x)
-        if self.expand_ratio == 1:
-            x = self.act(x)
+        x = self.bn2(x)
+        if self.expand_ratio == 1 and self.projection_activation:
+            x = self.projection_act(x)
 
         # Residual:
-        if self.strides == 1 and self.input_filters == self.output_filters:
+        if (
+            self.strides == 1
+            and self.input_filters == self.output_filters
+            and not self.nores
+        ):
             if self.dropout:
                 x = self.dropout_layer(x)
             x = keras.layers.Add(name=self.name + "add")([x, inputs])
@@ -205,10 +245,15 @@ class FusedMBConvBlock(keras.layers.Layer):
             "expand_ratio": self.expand_ratio,
             "kernel_size": self.kernel_size,
             "strides": self.strides,
+            "data_format": self.data_format,
             "se_ratio": self.se_ratio,
             "batch_norm_momentum": self.batch_norm_momentum,
+            "batch_norm_epsilon": self.batch_norm_epsilon,
             "activation": self.activation,
+            "projection_activation": self.projection_activation,
             "dropout": self.dropout,
+            "nores": self.nores,
+            "projection_kernel_size": self.projection_kernel_size,
         }
 
         base_config = super().get_config()

keras_hub/src/models/efficientnet/mbconv.py

@@ -2,15 +2,6 @@ import keras
 
 BN_AXIS = 3
 
-CONV_KERNEL_INITIALIZER = {
-    "class_name": "VarianceScaling",
-    "config": {
-        "scale": 2.0,
-        "mode": "fan_out",
-        "distribution": "truncated_normal",
-    },
-}
-
 
 class MBConvBlock(keras.layers.Layer):
     def __init__(
@@ -20,10 +11,13 @@ class MBConvBlock(keras.layers.Layer):
         expand_ratio=1,
         kernel_size=3,
         strides=1,
+        data_format="channels_last",
         se_ratio=0.0,
         batch_norm_momentum=0.9,
+        batch_norm_epsilon=1e-3,
         activation="swish",
         dropout=0.2,
+        nores=False,
         **kwargs
     ):
         """Implementation of the MBConv block
@@ -59,6 +53,9 @@ class MBConvBlock(keras.layers.Layer):
                 is above 0. The filters used in this phase are chosen as the
                 maximum between 1 and input_filters*se_ratio
             batch_norm_momentum: default 0.9, the BatchNormalization momentum
+            batch_norm_epsilon: default 1e-3, float, epsilon for batch norm
+                calcualtions. Used in denominator for calculations to prevent
+                divide by 0 errors.
             activation: default "swish", the activation function used between
                 convolution operations
             dropout: float, the optional dropout rate to apply before the output
@@ -79,10 +76,13 @@ class MBConvBlock(keras.layers.Layer):
         self.expand_ratio = expand_ratio
         self.kernel_size = kernel_size
         self.strides = strides
+        self.data_format = data_format
         self.se_ratio = se_ratio
         self.batch_norm_momentum = batch_norm_momentum
+        self.batch_norm_epsilon = batch_norm_epsilon
         self.activation = activation
         self.dropout = dropout
+        self.nores = nores
         self.filters = self.input_filters * self.expand_ratio
         self.filters_se = max(1, int(input_filters * se_ratio))
 
@@ -90,15 +90,16 @@ class MBConvBlock(keras.layers.Layer):
             filters=self.filters,
             kernel_size=1,
             strides=1,
-            kernel_initializer=CONV_KERNEL_INITIALIZER,
+            kernel_initializer=self._conv_kernel_initializer(),
             padding="same",
-            data_format="channels_last",
+            data_format=data_format,
             use_bias=False,
             name=self.name + "expand_conv",
         )
         self.bn1 = keras.layers.BatchNormalization(
             axis=BN_AXIS,
             momentum=self.batch_norm_momentum,
+            epsilon=self.batch_norm_epsilon,
             name=self.name + "expand_bn",
         )
         self.act = keras.layers.Activation(
@@ -107,9 +108,9 @@ class MBConvBlock(keras.layers.Layer):
         self.depthwise = keras.layers.DepthwiseConv2D(
             kernel_size=self.kernel_size,
             strides=self.strides,
-            depthwise_initializer=CONV_KERNEL_INITIALIZER,
+            depthwise_initializer=self._conv_kernel_initializer(),
             padding="same",
-            data_format="channels_last",
+            data_format=data_format,
             use_bias=False,
             name=self.name + "dwconv2",
         )
@@ -117,6 +118,7 @@ class MBConvBlock(keras.layers.Layer):
         self.bn2 = keras.layers.BatchNormalization(
             axis=BN_AXIS,
             momentum=self.batch_norm_momentum,
+            epsilon=self.batch_norm_epsilon,
             name=self.name + "bn",
         )
 
@@ -124,8 +126,9 @@ class MBConvBlock(keras.layers.Layer):
             self.filters_se,
             1,
             padding="same",
+            data_format=data_format,
             activation=self.activation,
-            kernel_initializer=CONV_KERNEL_INITIALIZER,
+            kernel_initializer=self._conv_kernel_initializer(),
             name=self.name + "se_reduce",
         )
 
@@ -133,18 +136,25 @@ class MBConvBlock(keras.layers.Layer):
             self.filters,
             1,
             padding="same",
+            data_format=data_format,
             activation="sigmoid",
-            kernel_initializer=CONV_KERNEL_INITIALIZER,
+            kernel_initializer=self._conv_kernel_initializer(),
             name=self.name + "se_expand",
         )
 
+        projection_kernel_size = 1 if expand_ratio != 1 else kernel_size
+        padding_pixels = projection_kernel_size // 2
+        self.output_conv_pad = keras.layers.ZeroPadding2D(
+            padding=(padding_pixels, padding_pixels),
+            name=self.name + "project_conv_pad",
+        )
         self.output_conv = keras.layers.Conv2D(
             filters=self.output_filters,
-            kernel_size=1 if expand_ratio != 1 else kernel_size,
+            kernel_size=projection_kernel_size,
             strides=1,
-            kernel_initializer=CONV_KERNEL_INITIALIZER,
-            padding="same",
-            data_format="channels_last",
+            kernel_initializer=self._conv_kernel_initializer(),
+            padding="valid",
+            data_format=data_format,
             use_bias=False,
             name=self.name + "project_conv",
         )
@@ -152,6 +162,7 @@ class MBConvBlock(keras.layers.Layer):
         self.bn3 = keras.layers.BatchNormalization(
             axis=BN_AXIS,
             momentum=self.batch_norm_momentum,
+            epsilon=self.batch_norm_epsilon,
             name=self.name + "project_bn",
         )
 
@@ -162,6 +173,17 @@ class MBConvBlock(keras.layers.Layer):
                 name=self.name + "drop",
             )
 
+    def _conv_kernel_initializer(
+        self,
+        scale=2.0,
+        mode="fan_out",
+        distribution="truncated_normal",
+        seed=None,
+    ):
+        return keras.initializers.VarianceScaling(
+            scale=scale, mode=mode, distribution=distribution, seed=seed
+        )
+
     def build(self, input_shape):
         if self.name is None:
             self.name = keras.backend.get_uid("block0")
@@ -183,7 +205,8 @@ class MBConvBlock(keras.layers.Layer):
         # Squeeze and excite
         if 0 < self.se_ratio <= 1:
             se = keras.layers.GlobalAveragePooling2D(
-                name=self.name + "se_squeeze"
+                name=self.name + "se_squeeze",
+                data_format=self.data_format,
             )(x)
             if BN_AXIS == 1:
                 se_shape = (self.filters, 1, 1)
@@ -199,10 +222,15 @@ class MBConvBlock(keras.layers.Layer):
             x = keras.layers.multiply([x, se], name=self.name + "se_excite")
 
         # Output phase
+        x = self.output_conv_pad(x)
         x = self.output_conv(x)
         x = self.bn3(x)
 
-        if self.strides == 1 and self.input_filters == self.output_filters:
+        if (
+            self.strides == 1
+            and self.input_filters == self.output_filters
+            and not self.nores
+        ):
             if self.dropout:
                 x = self.dropout_layer(x)
             x = keras.layers.Add(name=self.name + "add")([x, inputs])
@@ -215,10 +243,13 @@ class MBConvBlock(keras.layers.Layer):
             "expand_ratio": self.expand_ratio,
             "kernel_size": self.kernel_size,
             "strides": self.strides,
+            "data_format": self.data_format,
             "se_ratio": self.se_ratio,
             "batch_norm_momentum": self.batch_norm_momentum,
+            "batch_norm_epsilon": self.batch_norm_epsilon,
             "activation": self.activation,
             "dropout": self.dropout,
+            "nores": self.nores,
         }
         base_config = super().get_config()
         return dict(list(base_config.items()) + list(config.items()))