Repository: james77777778/keras-image-models (language: python)

Issue:
Add popular models
Reference:
- https://huggingface.co/timm?sort_models=downloads#models
- https://huggingface.co/collections/timm/timm-takes-on-the-classics-655ce63dcaa21067cd7d3f4c

Classics:
- [x] densenet
- [x] inceptionv3
- [ ] inceptionresnetv2
- [x] resnet
- [ ] squeezenet
- [x] vgg
- [x] xception

Lightweight
- [ ] edgenext
- [x] efficientnetlite
- [ ] efficientvit
- [x] ghostnet
- [x] ghostnetv2
- [x] ghostnetv3
- [x] lcnet
- [ ] levit
- [x] mobilenetv2
- [x] mobilenetv3
- [x] mobileone
- [x] mobilevit
- [x] mobilevitv2
- [ ] repghostnet
- [x] tinynet

SOTA:
- [x] convmixer
- [x] convnext
- [x] efficientnet
- [x] efficientnetv2
- [ ] eva02
- [x] hgnet
- [x] hgnetv2
- [x] inceptionnext
- [ ] maxvit
- [x] regnet
- [x] repvgg
- [ ] swin transformer
- [x] vit


Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index 9b0bbd2..6b51b06 100644
--- a/README.md
+++ b/README.md
@@ -10,7 +10,5 @@ CUDA_VISIBLE_DEVICES= KERAS_BACKEND=tensorflow pytest
 ## Work in Progress
 
 - Test pretrained weights
-- Add `MobileNet100V3SmallMinimal` pretrained weights
-- Add `MobileNet100V3LargeMinimal` pretrained weights
 
 ## Acknowledgments
diff --git a/kimm/blocks/base_block.py b/kimm/blocks/base_block.py
index c38c897..2a428c6 100644
--- a/kimm/blocks/base_block.py
+++ b/kimm/blocks/base_block.py
@@ -24,20 +24,25 @@ def apply_conv2d_block(
     groups=1,
     activation=None,
     use_depthwise=False,
+    add_skip=False,
     bn_momentum=0.9,
     bn_epsilon=1e-5,
+    padding=None,
     name="conv2d_block",
 ):
     if kernel_size is None:
         raise ValueError(
             f"kernel_size must be passed. Received: kernel_size={kernel_size}"
         )
+    input_channels = inputs.shape[-1]
+    has_skip = add_skip and strides == 1 and input_channels == filters
     x = inputs
 
-    padding = "same"
-    if strides > 1:
-        padding = "valid"
-        x = layers.ZeroPadding2D(kernel_size // 2, name=f"{name}_pad")(x)
+    if padding is None:
+        padding = "same"
+        if strides > 1:
+            padding = "valid"
+            x = layers.ZeroPadding2D(kernel_size // 2, name=f"{name}_pad")(x)
 
     if not use_depthwise:
         x = layers.Conv2D(
@@ -61,6 +66,8 @@ def apply_conv2d_block(
         name=f"{name}_bn", momentum=bn_momentum, epsilon=bn_epsilon
     )(x)
     x = apply_activation(x, activation, name=name)
+    if has_skip:
+        x = layers.Add()([x, inputs])
     return x
 
 
@@ -70,14 +77,17 @@ def apply_se_block(
     activation="relu",
     gate_activation="sigmoid",
     make_divisible_number=None,
+    se_input_channels=None,
     name="se_block",
 ):
     input_channels = inputs.shape[-1]
+    if se_input_channels is None:
+        se_input_channels = input_channels
     if make_divisible_number is None:
-        se_channels = round(input_channels * se_ratio)
+        se_channels = round(se_input_channels * se_ratio)
     else:
         se_channels = make_divisible(
-            input_channels * se_ratio, make_divisible_number
+            se_input_channels * se_ratio, make_divisible_number
         )
 
     ori_x = inputs
diff --git a/kimm/layers/attention.py b/kimm/layers/attention.py
index 5627e2b..1797020 100644
--- a/kimm/layers/attention.py
+++ b/kimm/layers/attention.py
@@ -26,23 +26,53 @@ def __init__(
         self.name = name
 
         self.qkv = layers.Dense(
-            hidden_dim * 3, use_bias=use_qkv_bias, name=f"{name}_qkv"
+            hidden_dim * 3,
+            use_bias=use_qkv_bias,
+            dtype=self.dtype,
+            name=f"{name}_qkv",
         )
         if use_qk_norm:
-            self.q_norm = layers.LayerNormalization(name=f"{name}_q_norm")
-            self.k_norm = layers.LayerNormalization(name=f"{name}_k_norm")
+            self.q_norm = layers.LayerNormalization(
+                dtype=self.dtype_policy, name=f"{name}_q_norm"
+            )
+            self.k_norm = layers.LayerNormalization(
+                dtype=self.dtype_policy, name=f"{name}_k_norm"
+            )
         else:
-            self.q_norm = layers.Identity()
-            self.k_norm = layers.Identity()
+            self.q_norm = layers.Identity(dtype=self.dtype_policy)
+            self.k_norm = layers.Identity(dtype=self.dtype_policy)
 
         self.attention_dropout = layers.Dropout(
-            attention_dropout_rate, name=f"{name}_attn_drop"
+            attention_dropout_rate,
+            dtype=self.dtype_policy,
+            name=f"{name}_attn_drop",
+        )
+        self.projection = layers.Dense(
+            hidden_dim, dtype=self.dtype_policy, name=f"{name}_proj"
         )
-        self.projection = layers.Dense(hidden_dim, name=f"{name}_proj")
         self.projection_dropout = layers.Dropout(
-            projection_dropout_rate, name=f"{name}_proj_drop"
+            projection_dropout_rate,
+            dtype=self.dtype_policy,
+            name=f"{name}_proj_drop",
         )
 
+    def build(self, input_shape):
+        self.qkv.build(input_shape)
+        qkv_output_shape = list(input_shape)
+        qkv_output_shape[-1] = qkv_output_shape[-1] * 3
+        self.q_norm.build(qkv_output_shape)
+        self.k_norm.build(qkv_output_shape)
+        attention_input_shape = [
+            input_shape[0],
+            self.num_heads,
+            input_shape[1],
+            input_shape[1],
+        ]
+        self.attention_dropout.build(attention_input_shape)
+        self.projection.build(input_shape)
+        self.projection_dropout.build(input_shape)
+        self.built = True
+
     def call(self, inputs, training=None, mask=None):
         input_shape = ops.shape(inputs)
         qkv = self.qkv(inputs)
diff --git a/kimm/layers/layer_scale.py b/kimm/layers/layer_scale.py
index 98620a9..9f030c7 100644
--- a/kimm/layers/layer_scale.py
+++ b/kimm/layers/layer_scale.py
@@ -30,7 +30,7 @@ def get_config(self):
         config.update(
             {
                 "hidden_size": self.hidden_size,
-                "initializer": self.initializer,
+                "initializer": initializers.serialize(self.initializer),
                 "name": self.name,
             }
         )
diff --git a/kimm/layers/position_embedding.py b/kimm/layers/position_embedding.py
index 4cb1643..167b9ae 100644
--- a/kimm/layers/position_embedding.py
+++ b/kimm/layers/position_embedding.py
@@ -7,6 +7,11 @@ def __init__(self, **kwargs):
         super().__init__(**kwargs)
 
     def build(self, input_shape):
+        if len(input_shape) != 3:
+            raise ValueError(
+                "PositionEmbedding only accepts 3-dimensional input. "
+                f"Received: input_shape={input_shape}"
+            )
         self.pos_embed = self.add_weight(
             shape=[1, input_shape[-2] + 1, input_shape[-1]],
             initializer="random_normal",
@@ -26,6 +31,11 @@ def call(self, inputs, training=None, mask=None):
         x = ops.add(x, self.pos_embed)
         return x
 
+    def compute_output_shape(self, input_shape):
+        output_shape = list(input_shape)
+        output_shape[1] = output_shape[1] + 1
+        return output_shape
+
     def get_config(self):
         return super().get_config()
 
diff --git a/kimm/models/efficientnet.py b/kimm/models/efficientnet.py
new file mode 100644
index 0000000..0ef966e
--- /dev/null
+++ b/kimm/models/efficientnet.py
@@ -0,0 +1,1430 @@
+import math
+import typing
+
+import keras
+from keras import backend
+from keras import layers
+from keras import utils
+from keras.src.applications import imagenet_utils
+
+from kimm.blocks import apply_conv2d_block
+from kimm.blocks import apply_se_block
+from kimm.models.feature_extractor import FeatureExtractor
+from kimm.utils import make_divisible
+
+DEFAULT_V1_CONFIG = [
+    # type, repeat, kernel_size, strides, expansion_ratio, channels, se_ratio
+    ["ds", 1, 3, 1, 1, 16, 0.25],
+    ["ir", 2, 3, 2, 6, 24, 0.25],
+    ["ir", 2, 5, 2, 6, 40, 0.25],
+    ["ir", 3, 3, 2, 6, 80, 0.25],
+    ["ir", 3, 5, 1, 6, 112, 0.25],
+    ["ir", 4, 5, 2, 6, 192, 0.25],
+    ["ir", 1, 3, 1, 6, 320, 0.25],
+]
+DEFAULT_V1_LITE_CONFIG = [
+    # type, repeat, kernel_size, strides, expansion_ratio, channels, se_ratio
+    ["ds", 1, 3, 1, 1, 16, 0.0],
+    ["ir", 2, 3, 2, 6, 24, 0.0],
+    ["ir", 2, 5, 2, 6, 40, 0.0],
+    ["ir", 3, 3, 2, 6, 80, 0.0],
+    ["ir", 3, 5, 1, 6, 112, 0.0],
+    ["ir", 4, 5, 2, 6, 192, 0.0],
+    ["ir", 1, 3, 1, 6, 320, 0.0],
+]
+DEFAULT_V2_S_CONFIG = [
+    # type, repeat, kernel_size, strides, expansion_ratio, channels, se_ratio
+    ["cn", 2, 3, 1, 1, 24, 0.0],
+    ["er", 4, 3, 2, 4, 48, 0.0],
+    ["er", 4, 3, 2, 4, 64, 0.0],
+    ["ir", 6, 3, 2, 4, 128, 0.25],
+    ["ir", 9, 3, 1, 6, 160, 0.25],
+    ["ir", 15, 3, 2, 6, 256, 0.25],
+]
+DEFAULT_V2_M_CONFIG = [
+    # type, repeat, kernel_size, strides, expansion_ratio, channels, se_ratio
+    ["cn", 3, 3, 1, 1, 24, 0.0],
+    ["er", 5, 3, 2, 4, 48, 0.0],
+    ["er", 5, 3, 2, 4, 80, 0.0],
+    ["ir", 7, 3, 2, 4, 160, 0.25],
+    ["ir", 14, 3, 1, 6, 176, 0.25],
+    ["ir", 18, 3, 2, 6, 304, 0.25],
+    ["ir", 5, 3, 1, 6, 512, 0.25],
+]
+DEFAULT_V2_L_CONFIG = [
+    # type, repeat, kernel_size, strides, expansion_ratio, channels, se_ratio
+    ["cn", 4, 3, 1, 1, 32, 0.0],
+    ["er", 7, 3, 2, 4, 64, 0.0],
+    ["er", 7, 3, 2, 4, 96, 0.0],
+    ["ir", 10, 3, 2, 4, 192, 0.25],
+    ["ir", 19, 3, 1, 6, 224, 0.25],
+    ["ir", 25, 3, 2, 6, 384, 0.25],
+    ["ir", 7, 3, 1, 6, 640, 0.25],
+]
+DEFAULT_V2_XL_CONFIG = [
+    # type, repeat, kernel_size, strides, expansion_ratio, channels, se_ratio
+    ["cn", 4, 3, 1, 1, 32, 0.0],
+    ["er", 8, 3, 2, 4, 64, 0.0],
+    ["er", 8, 3, 2, 4, 96, 0.0],
+    ["ir", 16, 3, 2, 4, 192, 0.25],
+    ["ir", 24, 3, 1, 6, 256, 0.25],
+    ["ir", 32, 3, 2, 6, 512, 0.25],
+    ["ir", 8, 3, 1, 6, 640, 0.25],
+]
+DEFAULT_V2_BASE_CONFIG = [
+    # type, repeat, kernel_size, strides, expansion_ratio, channels, se_ratio
+    ["cn", 1, 3, 1, 1, 16, 0.0],
+    ["er", 2, 3, 2, 4, 32, 0.0],
+    ["er", 2, 3, 2, 4, 48, 0.0],
+    ["ir", 3, 3, 2, 4, 96, 0.25],
+    ["ir", 5, 3, 1, 6, 112, 0.25],
+    ["ir", 8, 3, 2, 6, 192, 0.25],
+]
+
+
+def apply_depthwise_separation_block(
+    inputs,
+    output_channels,
+    depthwise_kernel_size=3,
+    pointwise_kernel_size=1,
+    strides=1,
+    se_ratio=0.0,
+    activation="swish",
+    bn_epsilon=1e-5,
+    padding=None,
+    name="depthwise_separation_block",
+):
+    input_channels = inputs.shape[-1]
+    has_skip = strides == 1 and input_channels == output_channels
+
+    x = inputs
+    x = apply_conv2d_block(
+        x,
+        kernel_size=depthwise_kernel_size,
+        strides=strides,
+        activation=activation,
+        use_depthwise=True,
+        bn_epsilon=bn_epsilon,
+        padding=padding,
+        name=f"{name}_conv_dw",
+    )
+    if se_ratio > 0:
+        x = apply_se_block(
+            x,
+            se_ratio,
+            activation=activation,
+            gate_activation="sigmoid",
+            name=f"{name}_se",
+        )
+    x = apply_conv2d_block(
+        x,
+        output_channels,
+        pointwise_kernel_size,
+        1,
+        activation=None,
+        bn_epsilon=bn_epsilon,
+        padding=padding,
+        name=f"{name}_conv_pw",
+    )
+    if has_skip:
+        x = layers.Add()([x, inputs])
+    return x
+
+
+def apply_inverted_residual_block(
+    inputs,
+    output_channels,
+    depthwise_kernel_size=3,
+    expansion_kernel_size=1,
+    pointwise_kernel_size=1,
+    strides=1,
+    expansion_ratio=1.0,
+    se_ratio=0.0,
+    activation="swish",
+    bn_epsilon=1e-5,
+    padding=None,
+    name="inverted_residual_block",
+):
+    input_channels = inputs.shape[-1]
+    hidden_channels = make_divisible(input_channels * expansion_ratio)
+    has_skip = strides == 1 and input_channels == output_channels
+
+    x = inputs
+
+    # Point-wise expansion
+    x = apply_conv2d_block(
+        x,
+        hidden_channels,
+        expansion_kernel_size,
+        1,
+        activation=activation,
+        bn_epsilon=bn_epsilon,
+        padding=padding,
+        name=f"{name}_conv_pw",
+    )
+    # Depth-wise convolution
+    x = apply_conv2d_block(
+        x,
+        kernel_size=depthwise_kernel_size,
+        strides=strides,
+        activation=activation,
+        use_depthwise=True,
+        bn_epsilon=bn_epsilon,
+        padding=padding,
+        name=f"{name}_conv_dw",
+    )
+    # Squeeze-and-excitation
+    if se_ratio > 0:
+        x = apply_se_block(
+            x,
+            se_ratio,
+            activation=activation,
+            gate_activation="sigmoid",
+            se_input_channels=input_channels,
+            name=f"{name}_se",
+        )
+    # Point-wise linear projection
+    x = apply_conv2d_block(
+        x,
+        output_channels,
+        pointwise_kernel_size,
+        1,
+        activation=None,
+        bn_epsilon=bn_epsilon,
+        padding=padding,
+        name=f"{name}_conv_pwl",
+    )
+    if has_skip:
+        x = layers.Add()([x, inputs])
+    return x
+
+
+def apply_edge_residual_block(
+    inputs,
+    output_channels,
+    expansion_kernel_size=1,
+    pointwise_kernel_size=1,
+    strides=1,
+    expansion_ratio=1.0,
+    se_ratio=0.0,
+    activation="swish",
+    bn_epsilon=1e-5,
+    padding=None,
+    name="edge_residual_block",
+):
+    input_channels = inputs.shape[-1]
+    hidden_channels = make_divisible(input_channels * expansion_ratio)
+    has_skip = strides == 1 and input_channels == output_channels
+
+    x = inputs
+
+    # Expansion
+    x = apply_conv2d_block(
+        x,
+        hidden_channels,
+        expansion_kernel_size,
+        strides,
+        activation=activation,
+        bn_epsilon=bn_epsilon,
+        padding=padding,
+        name=f"{name}_conv_exp",
+    )
+    # Squeeze-and-excitation
+    if se_ratio > 0:
+        x = apply_se_block(
+            x,
+            se_ratio,
+            activation=activation,
+            gate_activation="sigmoid",
+            se_input_channels=input_channels,
+            name=f"{name}_se",
+        )
+    # Point-wise linear projection
+    x = apply_conv2d_block(
+        x,
+        output_channels,
+        pointwise_kernel_size,
+        1,
+        activation=None,
+        bn_epsilon=bn_epsilon,
+        padding=padding,
+        name=f"{name}_conv_pwl",
+    )
+    if has_skip:
+        x = layers.Add()([x, inputs])
+    return x
+
+
+class EfficientNet(FeatureExtractor):
+    def __init__(
+        self,
+        width: float = 1.0,
+        depth: float = 1.0,
+        stem_channels: int = 32,
+        head_channels: int = 1280,
+        fix_stem_and_head_channels: bool = False,
+        fix_first_and_last_blocks: bool = False,
+        activation="swish",
+        input_tensor: keras.KerasTensor = None,
+        input_shape: typing.Optional[typing.Sequence[int]] = None,
+        include_preprocessing: bool = True,
+        include_top: bool = True,
+        pooling: typing.Optional[str] = None,
+        dropout_rate: float = 0.0,
+        classes: int = 1000,
+        classifier_activation: str = "softmax",
+        weights: typing.Optional[str] = None,  # TODO: imagenet
+        config: typing.Union[str, typing.List] = "default",
+        **kwargs,
+    ):
+        available_configs = [
+            "v1",
+            "v1_lite",
+            "v2_s",
+            "v2_m",
+            "v2_l",
+            "v2_xl",
+            "v2_base",
+        ]
+        if config == "v1":
+            config = DEFAULT_V1_CONFIG
+        elif config == "v1_lite":
+            config = DEFAULT_V1_LITE_CONFIG
+        elif config == "v2_s":
+            config = DEFAULT_V2_S_CONFIG
+        elif config == "v2_m":
+            config = DEFAULT_V2_M_CONFIG
+        elif config == "v2_l":
+            config = DEFAULT_V2_L_CONFIG
+        elif config == "v2_xl":
+            config = DEFAULT_V2_XL_CONFIG
+        elif config == "v2_base":
+            config = DEFAULT_V2_BASE_CONFIG
+        else:
+            if isinstance(config, str):
+                raise ValueError(
+                    f"config must be one of {available_configs} using string. "
+                    f"Received: config={config}"
+                )
+        # TF default config
+        default_size = kwargs.pop("default_size", 224)
+        bn_epsilon = kwargs.pop("bn_epsilon", 1e-5)
+        padding = kwargs.pop("padding", None)
+
+        # Prepare feature extraction
+        features = {}
+
+        # Determine proper input shape
+        input_shape = imagenet_utils.obtain_input_shape(
+            input_shape,
+            default_size=default_size,
+            min_size=32,
+            data_format=backend.image_data_format(),
+            require_flatten=include_top,
+            weights=weights,
+        )
+
+        if input_tensor is None:
+            img_input = layers.Input(shape=input_shape)
+        else:
+            if not backend.is_keras_tensor(input_tensor):
+                img_input = layers.Input(tensor=input_tensor, shape=input_shape)
+            else:
+                img_input = input_tensor
+
+        x = img_input
+
+        # [0, 255] to [0, 1] and apply ImageNet mean and variance
+        if include_preprocessing:
+            x = layers.Rescaling(scale=1.0 / 255.0)(x)
+            x = layers.Normalization(
+                mean=[0.485, 0.456, 0.406], variance=[0.229, 0.224, 0.225]
+            )(x)
+
+        # stem
+        stem_channel = (
+            stem_channels
+            if fix_stem_and_head_channels
+            else make_divisible(stem_channels * width)
+        )
+        x = apply_conv2d_block(
+            x,
+            stem_channel,
+            3,
+            2,
+            activation=activation,
+            bn_epsilon=bn_epsilon,
+            padding=padding,
+            name="conv_stem",
+        )
+        features["STEM_S2"] = x
+
+        # blocks
+        current_stride = 2
+        for current_block_idx, cfg in enumerate(config):
+            block_type, r, k, s, e, c, se = cfg
+            c = make_divisible(c * width)
+            if fix_first_and_last_blocks and (
+                current_block_idx in (0, len(config) - 1)
+            ):
+                r = r
+            else:
+                r = int(math.ceil(r * depth))
+            for current_layer_idx in range(r):
+                s = s if current_layer_idx == 0 else 1
+                common_kwargs = {
+                    "bn_epsilon": bn_epsilon,
+                    "padding": padding,
+                    "name": f"blocks_{current_block_idx}_{current_layer_idx}",
+                }
+                if block_type == "ds":
+                    x = apply_depthwise_separation_block(
+                        x, c, k, 1, s, se, activation, **common_kwargs
+                    )
+                elif block_type == "ir":
+                    x = apply_inverted_residual_block(
+                        x, c, k, 1, 1, s, e, se, activation, **common_kwargs
+                    )
+                elif block_type == "cn":
+                    x = apply_conv2d_block(
+                        x,
+                        filters=c,
+                        kernel_size=k,
+                        strides=s,
+                        activation=activation,
+                        add_skip=True,
+                        **common_kwargs,
+                    )
+                elif block_type == "er":
+                    x = apply_edge_residual_block(
+                        x, c, k, 1, s, e, se, activation, **common_kwargs
+                    )
+                current_stride *= s
+            features[f"BLOCK{current_block_idx}_S{current_stride}"] = x
+
+        # last conv
+        if fix_stem_and_head_channels:
+            conv_head_channels = head_channels
+        else:
+            conv_head_channels = make_divisible(head_channels * width)
+        x = apply_conv2d_block(
+            x,
+            conv_head_channels,
+            1,
+            1,
+            activation=activation,
+            bn_epsilon=bn_epsilon,
+            name="conv_head",
+        )
+
+        if include_top:
+            x = layers.GlobalAveragePooling2D(name="avg_pool")(x)
+            x = layers.Dropout(rate=dropout_rate, name="conv_head_dropout")(x)
+            x = layers.Dense(
+                classes, activation=classifier_activation, name="classifier"
+            )(x)
+        else:
+            if pooling == "avg":
+                x = layers.GlobalAveragePooling2D(name="avg_pool")(x)
+            elif pooling == "max":
+                x = layers.GlobalMaxPooling2D(name="max_pool")(x)
+
+        # Ensure that the model takes into account
+        # any potential predecessors of `input_tensor`.
+        if input_tensor is not None:
+            inputs = utils.get_source_inputs(input_tensor)
+        else:
+            inputs = img_input
+
+        super().__init__(inputs=inputs, outputs=x, features=features, **kwargs)
+
+        # All references to `self` below this line
+        self.width = width
+        self.depth = depth
+        self.fix_stem_and_head_channels = fix_stem_and_head_channels
+        self.include_preprocessing = include_preprocessing
+        self.include_top = include_top
+        self.pooling = pooling
+        self.dropout_rate = dropout_rate
+        self.classes = classes
+        self.classifier_activation = classifier_activation
+        self._weights = weights  # `self.weights` is been used internally
+        self.config = config
+
+    @staticmethod
+    def available_feature_keys():
+        feature_keys = ["STEM_S2"]
+        feature_keys.extend(
+            [
+                f"BLOCK{i}_S{j}"
+                for i, j in zip(range(7), [2, 4, 8, 16, 16, 32, 32])
+            ]
+        )
+        return feature_keys
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "width": self.width,
+                "input_shape": self.input_shape[1:],
+                "include_preprocessing": self.include_preprocessing,
+                "include_top": self.include_top,
+                "pooling": self.pooling,
+                "dropout_rate": self.dropout_rate,
+                "classes": self.classes,
+                "classifier_activation": self.classifier_activation,
+                "weights": self._weights,
+                "config": self.config,
+            }
+        )
+        return config
+
+
+"""
+Model Definition
+"""
+
+
+class EfficientNetB0(EfficientNet):
+    def __init__(
+        self,
+        input_tensor: keras.KerasTensor = None,
+        input_shape: typing.Optional[typing.Sequence[int]] = None,
+        include_preprocessing: bool = True,
+        include_top: bool = True,
+        pooling: typing.Optional[str] = None,
+        dropout_rate: float = 0.0,
+        classes: int = 1000,
+        classifier_activation: str = "softmax",
+        weights: typing.Optional[str] = None,  # TODO: imagenet
+        config: typing.Union[str, typing.List] = "v1",
+        name: str = "EfficientNetB0",
+        **kwargs,
+    ):
+        # default to TF configuration (bn_epsilon=1e-3 and padding="same")
+        super().__init__(
+            1.0,
+            1.0,
+            32,
+            1280,
+            False,
+            False,
+            "swish",
+            input_tensor,
+            input_shape,
+            include_preprocessing,
+            include_top,
+            pooling,
+            dropout_rate,
+            classes,
+            classifier_activation,
+            weights,
+            config,
+            name=name,
+            default_size=224,
+            bn_epsilon=1e-3,
+            padding="same",
+            **kwargs,
+        )
+
+
+class EfficientNetB1(EfficientNet):
+    def __init__(
+        self,
+        input_tensor: keras.KerasTensor = None,
+        input_shape: typing.Optional[typing.Sequence[int]] = None,
+        include_preprocessing: bool = True,
+        include_top: bool = True,
+        pooling: typing.Optional[str] = None,
+        dropout_rate: float = 0.0,
+        classes: int = 1000,
+        classifier_activation: str = "softmax",
+        weights: typing.Optional[str] = None,  # TODO: imagenet
+        config: typing.Union[str, typing.List] = "v1",
+        name: str = "EfficientNetB1",
+        **kwargs,
+    ):
+        # default to TF configuration (bn_epsilon=1e-3 and padding="same")
+        super().__init__(
+            1.0,
+            1.1,
+            32,
+            1280,
+            False,
+            False,
+            "swish",
+            input_tensor,
+            input_shape,
+            include_preprocessing,
+            include_top,
+            pooling,
+            dropout_rate,
+            classes,
+            classifier_activation,
+            weights,
+            config,
+            name=name,
+            default_size=240,
+            bn_epsilon=1e-3,
+            padding="same",
+            **kwargs,
+        )
+
+
+class EfficientNetB2(EfficientNet):
+    def __init__(
+        self,
+        input_tensor: keras.KerasTensor = None,
+        input_shape: typing.Optional[typing.Sequence[int]] = None,
+        include_preprocessing: bool = True,
+        include_top: bool = True,
+        pooling: typing.Optional[str] = None,
+        dropout_rate: float = 0.0,
+        classes: int = 1000,
+        classifier_activation: str = "softmax",
+        weights: typing.Optional[str] = None,  # TODO: imagenet
+        config: typing.Union[str, typing.List] = "v1",
+        name: str = "EfficientNetB2",
+        **kwargs,
+    ):
+        # default to TF configuration (bn_epsilon=1e-3 and padding="same")
+        super().__init__(
+            1.1,
+            1.2,
+            32,
+            1280,
+            False,
+            False,
+            "swish",
+            input_tensor,
+            input_shape,
+            include_preprocessing,
+            include_top,
+            pooling,
+            dropout_rate,
+            classes,
+            classifier_activation,
+            weights,
+            config,
+            name=name,
+            default_size=260,
+            bn_epsilon=1e-3,
+            padding="same",
+            **kwargs,
+        )
+
+
+class EfficientNetB3(EfficientNet):
+    def __init__(
+        self,
+        input_tensor: keras.KerasTensor = None,
+        input_shape: typing.Optional[typing.Sequence[int]] = None,
+        include_preprocessing: bool = True,
+        include_top: bool = True,
+        pooling: typing.Optional[str] = None,
+        dropout_rate: float = 0.0,
+        classes: int = 1000,
+        classifier_activation: str = "softmax",
+        weights: typing.Optional[str] = None,  # TODO: imagenet
+        config: typing.Union[str, typing.List] = "v1",
+        name: str = "EfficientNetB3",
+        **kwargs,
+    ):
+        # default to TF configuration (bn_epsilon=1e-3 and padding="same")
+        super().__init__(
+            1.2,
+            1.4,
+            32,
+            1280,
+            False,
+            False,
+            "swish",
+            input_tensor,
+            input_shape,
+            include_preprocessing,
+            include_top,
+            pooling,
+            dropout_rate,
+            classes,
+            classifier_activation,
+            weights,
+            config,
+            name=name,
+            default_size=300,
+            bn_epsilon=1e-3,
+            padding="same",
+            **kwargs,
+        )
+
+
+class EfficientNetB4(EfficientNet):
+    def __init__(
+        self,
+        input_tensor: keras.KerasTensor = None,
+        input_shape: typing.Optional[typing.Sequence[int]] = None,
+        include_preprocessing: bool = True,
+        include_top: bool = True,
+        pooling: typing.Optional[str] = None,
+        dropout_rate: float = 0.0,
+        classes: int = 1000,
+        classifier_activation: str = "softmax",
+        weights: typing.Optional[str] = None,  # TODO: imagenet
+        config: typing.Union[str, typing.List] = "v1",
+        name: str = "EfficientNetB4",
+        **kwargs,
+    ):
+        # default to TF configuration (bn_epsilon=1e-3 and padding="same")
+        super().__init__(
+            1.4,
+            1.8,
+            32,
+            1280,
+            False,
+            False,
+            "swish",
+            input_tensor,
+            input_shape,
+            include_preprocessing,
+            include_top,
+            pooling,
+            dropout_rate,
+            classes,
+            classifier_activation,
+            weights,
+            config,
+            name=name,
+            default_size=380,
+            bn_epsilon=1e-3,
+            padding="same",
+            **kwargs,
+        )
+
+
+class EfficientNetB5(EfficientNet):
+    def __init__(
+        self,
+        input_tensor: keras.KerasTensor = None,
+        input_shape: typing.Optional[typing.Sequence[int]] = None,
+        include_preprocessing: bool = True,
+        include_top: bool = True,
+        pooling: typing.Optional[str] = None,
+        dropout_rate: float = 0.0,
+        classes: int = 1000,
+        classifier_activation: str = "softmax",
+        weights: typing.Optional[str] = None,  # TODO: imagenet
+        config: typing.Union[str, typing.List] = "v1",
+        name: str = "EfficientNetB5",
+        **kwargs,
+    ):
+        # default to TF configuration (bn_epsilon=1e-3 and padding="same")
+        super().__init__(
+            1.6,
+            2.2,
+            32,
+            1280,
+            False,
+            False,
+            "swish",
+            input_tensor,
+            input_shape,
+            include_preprocessing,
+            include_top,
+            pooling,
+            dropout_rate,
+            classes,
+            classifier_activation,
+            weights,
+            config,
+            name=name,
+            default_size=456,
+            bn_epsilon=1e-3,
+            padding="same",
+            **kwargs,
+        )
+
+
+class EfficientNetB6(EfficientNet):
+    def __init__(
+        self,
+        input_tensor: keras.KerasTensor = None,
+        input_shape: typing.Optional[typing.Sequence[int]] = None,
+        include_preprocessing: bool = True,
+        include_top: bool = True,
+        pooling: typing.Optional[str] = None,
+        dropout_rate: float = 0.0,
+        classes: int = 1000,
+        classifier_activation: str = "softmax",
+        weights: typing.Optional[str] = None,  # TODO: imagenet
+        config: typing.Union[str, typing.List] = "v1",
+        name: str = "EfficientNetB6",
+        **kwargs,
+    ):
+        # default to TF configuration (bn_epsilon=1e-3 and padding="same")
+        super().__init__(
+            1.8,
+            2.6,
+            32,
+            1280,
+            False,
+            False,
+            "swish",
+            input_tensor,
+            input_shape,
+            include_preprocessing,
+            include_top,
+            pooling,
+            dropout_rate,
+            classes,
+            classifier_activation,
+            weights,
+            config,
+            name=name,
+            default_size=528,
+            bn_epsilon=1e-3,
+            padding="same",
+            **kwargs,
+        )
+
+
+class EfficientNetB7(EfficientNet):
+    def __init__(
+        self,
+        input_tensor: keras.KerasTensor = None,
+        input_shape: typing.Optional[typing.Sequence[int]] = None,
+        include_preprocessing: bool = True,
+        include_top: bool = True,
+        pooling: typing.Optional[str] = None,
+        dropout_rate: float = 0.0,
+        classes: int = 1000,
+        classifier_activation: str = "softmax",
+        weights: typing.Optional[str] = None,  # TODO: imagenet
+        config: typing.Union[str, typing.List] = "v1",
+        name: str = "EfficientNetB7",
+        **kwargs,
+    ):
+        # default to TF configuration (bn_epsilon=1e-3 and padding="same")
+        super().__init__(
+            2.0,
+            3.1,
+            32,
+            1280,
+            False,
+            False,
+            "swish",
+            input_tensor,
+            input_shape,
+            include_preprocessing,
+            include_top,
+            pooling,
+            dropout_rate,
+            classes,
+            classifier_activation,
+            weights,
+            config,
+            name=name,
+            default_size=600,
+            bn_epsilon=1e-3,
+            padding="same",
+            **kwargs,
+        )
+
+
+class EfficientNetLiteB0(EfficientNet):
+    def __init__(
+        self,
+        input_tensor: keras.KerasTensor = None,
+        input_shape: typing.Optional[typing.Sequence[int]] = None,
+        include_preprocessing: bool = True,
+        include_top: bool = True,
+        pooling: typing.Optional[str] = None,
+        dropout_rate: float = 0.0,
+        classes: int = 1000,
+        classifier_activation: str = "softmax",
+        weights: typing.Optional[str] = None,  # TODO: imagenet
+        config: typing.Union[str, typing.List] = "v1_lite",
+        name: str = "EfficientNetLiteB0",
+        **kwargs,
+    ):
+        # default to TF configuration (bn_epsilon=1e-3 and padding="same")
+        super().__init__(
+            1.0,
+            1.0,
+            32,
+            1280,
+            True,
+            True,
+            "relu6",
+            input_tensor,
+            input_shape,
+            include_preprocessing,
+            include_top,
+            pooling,
+            dropout_rate,
+            classes,
+            classifier_activation,
+            weights,
+            config,
+            name=name,
+            default_size=224,
+            bn_epsilon=1e-3,
+            padding="same",
+            **kwargs,
+        )
+
+
+class EfficientNetLiteB1(EfficientNet):
+    def __init__(
+        self,
+        input_tensor: keras.KerasTensor = None,
+        input_shape: typing.Optional[typing.Sequence[int]] = None,
+        include_preprocessing: bool = True,
+        include_top: bool = True,
+        pooling: typing.Optional[str] = None,
+        dropout_rate: float = 0.0,
+        classes: int = 1000,
+        classifier_activation: str = "softmax",
+        weights: typing.Optional[str] = None,  # TODO: imagenet
+        config: typing.Union[str, typing.List] = "v1_lite",
+        name: str = "EfficientNetLiteB1",
+        **kwargs,
+    ):
+        # default to TF configuration (bn_epsilon=1e-3 and padding="same")
+        super().__init__(
+            1.0,
+            1.1,
+            32,
+            1280,
+            True,
+            True,
+            "relu6",
+            input_tensor,
+            input_shape,
+            include_preprocessing,
+            include_top,
+            pooling,
+            dropout_rate,
+            classes,
+            classifier_activation,
+            weights,
+            config,
+            name=name,
+            default_size=240,
+            bn_epsilon=1e-3,
+            padding="same",
+            **kwargs,
+        )
+
+
+class EfficientNetLiteB2(EfficientNet):
+    def __init__(
+        self,
+        input_tensor: keras.KerasTensor = None,
+        input_shape: typing.Optional[typing.Sequence[int]] = None,
+        include_preprocessing: bool = True,
+        include_top: bool = True,
+        pooling: typing.Optional[str] = None,
+        dropout_rate: float = 0.0,
+        classes: int = 1000,
+        classifier_activation: str = "softmax",
+        weights: typing.Optional[str] = None,  # TODO: imagenet
+        config: typing.Union[str, typing.List] = "v1_lite",
+        name: str = "EfficientNetLiteB2",
+        **kwargs,
+    ):
+        # default to TF configuration (bn_epsilon=1e-3 and padding="same")
+        super().__init__(
+            1.1,
+            1.2,
+            32,
+            1280,
+            True,
+            True,
+            "relu6",
+            input_tensor,
+            input_shape,
+            include_preprocessing,
+            include_top,
+            pooling,
+            dropout_rate,
+            classes,
+            classifier_activation,
+            weights,
+            config,
+            name=name,
+            default_size=260,
+            bn_epsilon=1e-3,
+            padding="same",
+            **kwargs,
+        )
+
+
+class EfficientNetLiteB3(EfficientNet):
+    def __init__(
+        self,
+        input_tensor: keras.KerasTensor = None,
+        input_shape: typing.Optional[typing.Sequence[int]] = None,
+        include_preprocessing: bool = True,
+        include_top: bool = True,
+        pooling: typing.Optional[str] = None,
+        dropout_rate: float = 0.0,
+        classes: int = 1000,
+        classifier_activation: str = "softmax",
+        weights: typing.Optional[str] = None,  # TODO: imagenet
+        config: typing.Union[str, typing.List] = "v1_lite",
+        name: str = "EfficientNetLiteB3",
+        **kwargs,
+    ):
+        # default to TF configuration (bn_epsilon=1e-3 and padding="same")
+        super().__init__(
+            1.2,
+            1.4,
+            32,
+            1280,
+            True,
+            True,
+            "relu6",
+            input_tensor,
+            input_shape,
+            include_preprocessing,
+            include_top,
+            pooling,
+            dropout_rate,
+            classes,
+            classifier_activation,
+            weights,
+            config,
+            name=name,
+            default_size=300,
+            bn_epsilon=1e-3,
+            padding="same",
+            **kwargs,
+        )
+
+
+class EfficientNetLiteB4(EfficientNet):
+    def __init__(
+        self,
+        input_tensor: keras.KerasTensor = None,
+        input_shape: typing.Optional[typing.Sequence[int]] = None,
+        include_preprocessing: bool = True,
+        include_top: bool = True,
+        pooling: typing.Optional[str] = None,
+        dropout_rate: float = 0.0,
+        classes: int = 1000,
+        classifier_activation: str = "softmax",
+        weights: typing.Optional[str] = None,  # TODO: imagenet
+        config: typing.Union[str, typing.List] = "v1_lite",
+        name: str = "EfficientNetLiteB4",
+        **kwargs,
+    ):
+        # default to TF configuration (bn_epsilon=1e-3 and padding="same")
+        super().__init__(
+            1.4,
+            1.8,
+            32,
+            1280,
+            True,
+            True,
+            "relu6",
+            input_tensor,
+            input_shape,
+            include_preprocessing,
+            include_top,
+            pooling,
+            dropout_rate,
+            classes,
+            classifier_activation,
+            weights,
+            config,
+            name=name,
+            default_size=380,
+            bn_epsilon=1e-3,
+            padding="same",
+            **kwargs,
+        )
+
+
+class EfficientNetV2S(EfficientNet):
+    def __init__(
+        self,
+        input_tensor: keras.KerasTensor = None,
+        input_shape: typing.Optional[typing.Sequence[int]] = None,
+        include_preprocessing: bool = True,
+        include_top: bool = True,
+        pooling: typing.Optional[str] = None,
+        dropout_rate: float = 0.0,
+        classes: int = 1000,
+        classifier_activation: str = "softmax",
+        weights: typing.Optional[str] = None,  # TODO: imagenet
+        config: typing.Union[str, typing.List] = "v2_s",
+        name: str = "EfficientNetV2S",
+        **kwargs,
+    ):
+        # default to TF configuration (bn_epsilon=1e-3 and padding="same")
+        super().__init__(
+            1.0,
+            1.0,
+            24,
+            1280,
+            False,
+            False,
+            "swish",
+            input_tensor,
+            input_shape,
+            include_preprocessing,
+            include_top,
+            pooling,
+            dropout_rate,
+            classes,
+            classifier_activation,
+            weights,
+            config,
+            name=name,
+            default_size=300,
+            bn_epsilon=1e-3,
+            padding="same",
+            **kwargs,
+        )
+
+    @staticmethod
+    def available_feature_keys():
+        feature_keys = ["STEM_S2"]
+        feature_keys.extend(
+            [f"BLOCK{i}_S{j}" for i, j in zip(range(6), [2, 4, 8, 16, 16, 32])]
+        )
+        return feature_keys
+
+
+class EfficientNetV2M(EfficientNet):
+    def __init__(
+        self,
+        input_tensor: keras.KerasTensor = None,
+        input_shape: typing.Optional[typing.Sequence[int]] = None,
+        include_preprocessing: bool = True,
+        include_top: bool = True,
+        pooling: typing.Optional[str] = None,
+        dropout_rate: float = 0.0,
+        classes: int = 1000,
+        classifier_activation: str = "softmax",
+        weights: typing.Optional[str] = None,  # TODO: imagenet
+        config: typing.Union[str, typing.List] = "v2_m",
+        name: str = "EfficientNetV2M",
+        **kwargs,
+    ):
+        # default to TF configuration (bn_epsilon=1e-3 and padding="same")
+        super().__init__(
+            1.0,
+            1.0,
+            24,
+            1280,
+            False,
+            False,
+            "swish",
+            input_tensor,
+            input_shape,
+            include_preprocessing,
+            include_top,
+            pooling,
+            dropout_rate,
+            classes,
+            classifier_activation,
+            weights,
+            config,
+            name=name,
+            default_size=384,
+            bn_epsilon=1e-3,
+            padding="same",
+            **kwargs,
+        )
+
+
+class EfficientNetV2L(EfficientNet):
+    def __init__(
+        self,
+        input_tensor: keras.KerasTensor = None,
+        input_shape: typing.Optional[typing.Sequence[int]] = None,
+        include_preprocessing: bool = True,
+        include_top: bool = True,
+        pooling: typing.Optional[str] = None,
+        dropout_rate: float = 0.0,
+        classes: int = 1000,
+        classifier_activation: str = "softmax",
+        weights: typing.Optional[str] = None,  # TODO: imagenet
+        config: typing.Union[str, typing.List] = "v2_l",
+        name: str = "EfficientNetV2L",
+        **kwargs,
+    ):
+        # default to TF configuration (bn_epsilon=1e-3 and padding="same")
+        super().__init__(
+            1.0,
+            1.0,
+            32,
+            1280,
+            False,
+            False,
+            "swish",
+            input_tensor,
+            input_shape,
+            include_preprocessing,
+            include_top,
+            pooling,
+            dropout_rate,
+            classes,
+            classifier_activation,
+            weights,
+            config,
+            name=name,
+            default_size=384,
+            bn_epsilon=1e-3,
+            padding="same",
+            **kwargs,
+        )
+
+
+class EfficientNetV2XL(EfficientNet):
+    def __init__(
+        self,
+        input_tensor: keras.KerasTensor = None,
+        input_shape: typing.Optional[typing.Sequence[int]] = None,
+        include_preprocessing: bool = True,
+        include_top: bool = True,
+        pooling: typing.Optional[str] = None,
+        dropout_rate: float = 0.0,
+        classes: int = 1000,
+        classifier_activation: str = "softmax",
+        weights: typing.Optional[str] = None,  # TODO: imagenet
+        config: typing.Union[str, typing.List] = "v2_xl",
+        name: str = "EfficientNetV2XL",
+        **kwargs,
+    ):
+        # default to TF configuration (bn_epsilon=1e-3 and padding="same")
+        super().__init__(
+            1.0,
+            1.0,
+            32,
+            1280,
+            False,
+            False,
+            "swish",
+            input_tensor,
+            input_shape,
+            include_preprocessing,
+            include_top,
+            pooling,
+            dropout_rate,
+            classes,
+            classifier_activation,
+            weights,
+            config,
+            name=name,
+            default_size=384,
+            bn_epsilon=1e-3,
+            padding="same",
+            **kwargs,
+        )
+
+
+class EfficientNetV2B0(EfficientNet):
+    def __init__(
+        self,
+        input_tensor: keras.KerasTensor = None,
+        input_shape: typing.Optional[typing.Sequence[int]] = None,
+        include_preprocessing: bool = True,
+        include_top: bool = True,
+        pooling: typing.Optional[str] = None,
+        dropout_rate: float = 0.0,
+        classes: int = 1000,
+        classifier_activation: str = "softmax",
+        weights: typing.Optional[str] = None,  # TODO: imagenet
+        config: typing.Union[str, typing.List] = "v2_base",
+        name: str = "EfficientNetV2B0",
+        **kwargs,
+    ):
+        # default to TF configuration (bn_epsilon=1e-3 and padding="same")
+        super().__init__(
+            1.0,
+            1.0,
+            32,
+            1280,
+            True,
+            False,
+            "swish",
+            input_tensor,
+            input_shape,
+            include_preprocessing,
+            include_top,
+            pooling,
+            dropout_rate,
+            classes,
+            classifier_activation,
+            weights,
+            config,
+            name=name,
+            default_size=192,
+            bn_epsilon=1e-3,
+            padding="same",
+            **kwargs,
+        )
+
+    @staticmethod
+    def available_feature_keys():
+        feature_keys = ["STEM_S2"]
+        feature_keys.extend(
+            [f"BLOCK{i}_S{j}" for i, j in zip(range(6), [2, 4, 8, 16, 16, 32])]
+        )
+        return feature_keys
+
+
+class EfficientNetV2B1(EfficientNet):
+    def __init__(
+        self,
+        input_tensor: keras.KerasTensor = None,
+        input_shape: typing.Optional[typing.Sequence[int]] = None,
+        include_preprocessing: bool = True,
+        include_top: bool = True,
+        pooling: typing.Optional[str] = None,
+        dropout_rate: float = 0.0,
+        classes: int = 1000,
+        classifier_activation: str = "softmax",
+        weights: typing.Optional[str] = None,  # TODO: imagenet
+        config: typing.Union[str, typing.List] = "v2_base",
+        name: str = "EfficientNetV2B1",
+        **kwargs,
+    ):
+        # default to TF configuration (bn_epsilon=1e-3 and padding="same")
+        super().__init__(
+            1.0,
+            1.1,
+            32,
+            1280,
+            True,
+            False,
+            "swish",
+            input_tensor,
+            input_shape,
+            include_preprocessing,
+            include_top,
+            pooling,
+            dropout_rate,
+            classes,
+            classifier_activation,
+            weights,
+            config,
+            name=name,
+            default_size=192,
+            bn_epsilon=1e-3,
+            padding="same",
+            **kwargs,
+        )
+
+    @staticmethod
+    def available_feature_keys():
+        feature_keys = ["STEM_S2"]
+        feature_keys.extend(
+            [f"BLOCK{i}_S{j}" for i, j in zip(range(6), [2, 4, 8, 16, 16, 32])]
+        )
+        return feature_keys
+
+
+class EfficientNetV2B2(EfficientNet):
+    def __init__(
+        self,
+        input_tensor: keras.KerasTensor = None,
+        input_shape: typing.Optional[typing.Sequence[int]] = None,
+        include_preprocessing: bool = True,
+        include_top: bool = True,
+        pooling: typing.Optional[str] = None,
+        dropout_rate: float = 0.0,
+        classes: int = 1000,
+        classifier_activation: str = "softmax",
+        weights: typing.Optional[str] = None,  # TODO: imagenet
+        config: typing.Union[str, typing.List] = "v2_base",
+        name: str = "EfficientNetV2B2",
+        **kwargs,
+    ):
+        # default to TF configuration (bn_epsilon=1e-3 and padding="same")
+        super().__init__(
+            1.1,
+            1.2,
+            32,
+            make_divisible(1280 * 1.1),
+            True,
+            False,
+            "swish",
+            input_tensor,
+            input_shape,
+            include_preprocessing,
+            include_top,
+            pooling,
+            dropout_rate,
+            classes,
+            classifier_activation,
+            weights,
+            config,
+            name=name,
+            default_size=208,
+            bn_epsilon=1e-3,
+            padding="same",
+            **kwargs,
+        )
+
+    @staticmethod
+    def available_feature_keys():
+        feature_keys = ["STEM_S2"]
+        feature_keys.extend(
+            [f"BLOCK{i}_S{j}" for i, j in zip(range(6), [2, 4, 8, 16, 16, 32])]
+        )
+        return feature_keys
+
+
+# TODO: fix pretrained weights for EfficientNetV2B3
+class EfficientNetV2B3(EfficientNet):
+    def __init__(
+        self,
+        input_tensor: keras.KerasTensor = None,
+        input_shape: typing.Optional[typing.Sequence[int]] = None,
+        include_preprocessing: bool = True,
+        include_top: bool = True,
+        pooling: typing.Optional[str] = None,
+        dropout_rate: float = 0.0,
+        classes: int = 1000,
+        classifier_activation: str = "softmax",
+        weights: typing.Optional[str] = None,  # TODO: imagenet
+        config: typing.Union[str, typing.List] = "v2_base",
+        name: str = "EfficientNetV2B3",
+        **kwargs,
+    ):
+        # default to TF configuration (bn_epsilon=1e-3 and padding="same")
+        super().__init__(
+            1.2,
+            1.4,
+            32,
+            make_divisible(1280 * 1.2),
+            True,
+            False,
+            "swish",
+            input_tensor,
+            input_shape,
+            include_preprocessing,
+            include_top,
+            pooling,
+            dropout_rate,
+            classes,
+            classifier_activation,
+            weights,
+            config,
+            name=name,
+            default_size=240,
+            bn_epsilon=1e-3,
+            padding="same",
+            **kwargs,
+        )
+
+    @staticmethod
+    def available_feature_keys():
+        feature_keys = ["STEM_S2"]
+        feature_keys.extend(
+            [f"BLOCK{i}_S{j}" for i, j in zip(range(6), [2, 4, 8, 16, 16, 32])]
+        )
+        return feature_keys
diff --git a/kimm/models/ghostnet.py b/kimm/models/ghostnet.py
index 6d93757..b8a3231 100644
--- a/kimm/models/ghostnet.py
+++ b/kimm/models/ghostnet.py
@@ -414,15 +414,8 @@ def __init__(
         name: str = "GhostNet050",
         **kwargs,
     ):
-        expected_width = 0.5
-        width = kwargs.pop("width", expected_width)
-        if width != expected_width:
-            raise ValueError(
-                f"`width` must be {expected_width} for "
-                f"{self.__class__.__name__}. Got `width`={width}"
-            )
         super().__init__(
-            width,
+            0.5,
             input_tensor,
             input_shape,
             include_preprocessing,
@@ -454,15 +447,8 @@ def __init__(
         name: str = "GhostNet100",
         **kwargs,
     ):
-        expected_width = 1.0
-        width = kwargs.pop("width", expected_width)
-        if width != expected_width:
-            raise ValueError(
-                f"`width` must be {expected_width} for "
-                f"{self.__class__.__name__}. Got `width`={width}"
-            )
         super().__init__(
-            width,
+            1.0,
             input_tensor,
             input_shape,
             include_preprocessing,
@@ -494,15 +480,8 @@ def __init__(
         name: str = "GhostNet130",
         **kwargs,
     ):
-        expected_width = 1.3
-        width = kwargs.pop("width", expected_width)
-        if width != expected_width:
-            raise ValueError(
-                f"`width` must be {expected_width} for "
-                f"{self.__class__.__name__}. Got `width`={width}"
-            )
         super().__init__(
-            width,
+            1.3,
             input_tensor,
             input_shape,
             include_preprocessing,
@@ -534,15 +513,8 @@ def __init__(
         name: str = "GhostNet100V2",
         **kwargs,
     ):
-        expected_width = 1.0
-        width = kwargs.pop("width", expected_width)
-        if width != expected_width:
-            raise ValueError(
-                f"`width` must be {expected_width} for "
-                f"{self.__class__.__name__}. Got `width`={width}"
-            )
         super().__init__(
-            width,
+            1.0,
             input_tensor,
             input_shape,
             include_preprocessing,
@@ -575,15 +547,8 @@ def __init__(
         name: str = "GhostNet130V2",
         **kwargs,
     ):
-        expected_width = 1.3
-        width = kwargs.pop("width", expected_width)
-        if width != expected_width:
-            raise ValueError(
-                f"`width` must be {expected_width} for "
-                f"{self.__class__.__name__}. Got `width`={width}"
-            )
         super().__init__(
-            width,
+            1.3,
             input_tensor,
             input_shape,
             include_preprocessing,
@@ -616,15 +581,8 @@ def __init__(
         name: str = "GhostNet160V2",
         **kwargs,
     ):
-        expected_width = 1.6
-        width = kwargs.pop("width", expected_width)
-        if width != expected_width:
-            raise ValueError(
-                f"`width` must be {expected_width} for "
-                f"{self.__class__.__name__}. Got `width`={width}"
-            )
         super().__init__(
-            width,
+            1.6,
             input_tensor,
             input_shape,
             include_preprocessing,
diff --git a/kimm/models/mobilenet_v3.py b/kimm/models/mobilenet_v3.py
index 997d229..6eff96c 100644
--- a/kimm/models/mobilenet_v3.py
+++ b/kimm/models/mobilenet_v3.py
@@ -69,6 +69,8 @@ def apply_depthwise_separation_block(
     strides=1,
     se_ratio=0.0,
     activation="relu",
+    bn_epsilon=1e-5,
+    padding=None,
     name="depthwise_separation_block",
 ):
     input_channels = inputs.shape[-1]
@@ -81,6 +83,8 @@ def apply_depthwise_separation_block(
         strides=strides,
         activation=activation,
         use_depthwise=True,
+        bn_epsilon=bn_epsilon,
+        padding=padding,
         name=f"{name}_conv_dw",
     )
     if se_ratio > 0:
@@ -98,6 +102,8 @@ def apply_depthwise_separation_block(
         pointwise_kernel_size,
         1,
         activation=None,
+        bn_epsilon=bn_epsilon,
+        padding=padding,
         name=f"{name}_conv_pw",
     )
     if has_skip:
@@ -115,6 +121,8 @@ def apply_inverted_residual_block(
     expansion_ratio=1.0,
     se_ratio=0.0,
     activation="relu",
+    bn_epsilon=1e-5,
+    padding=None,
     name="inverted_residual_block",
 ):
     input_channels = inputs.shape[-1]
@@ -130,6 +138,8 @@ def apply_inverted_residual_block(
         expansion_kernel_size,
         1,
         activation=activation,
+        bn_epsilon=bn_epsilon,
+        padding=padding,
         name=f"{name}_conv_pw",
     )
     # Depth-wise convolution
@@ -139,6 +149,8 @@ def apply_inverted_residual_block(
         strides=strides,
         activation=activation,
         use_depthwise=True,
+        bn_epsilon=bn_epsilon,
+        padding=padding,
         name=f"{name}_conv_dw",
     )
     # Squeeze-and-excitation
@@ -158,6 +170,8 @@ def apply_inverted_residual_block(
         pointwise_kernel_size,
         1,
         activation=None,
+        bn_epsilon=bn_epsilon,
+        padding=padding,
         name=f"{name}_conv_pwl",
     )
     if has_skip:
@@ -198,6 +212,9 @@ def __init__(
             force_activation = None
             force_kernel_size = None
             no_se = False
+        # TF default config
+        bn_epsilon = kwargs.pop("bn_epsilon", 1e-5)
+        padding = kwargs.pop("padding", None)
 
         # Prepare feature extraction
         features = {}
@@ -239,6 +256,8 @@ def __init__(
             3,
             2,
             activation=force_activation or "hard_swish",
+            bn_epsilon=bn_epsilon,
+            padding=padding,
             name="conv_stem",
         )
         features["STEM_S2"] = x
@@ -263,17 +282,21 @@ def __init__(
                     r = int(math.ceil(r * depth))
                 for current_layer_idx in range(r):
                     s = s if current_layer_idx == 0 else 1
-                    name = (
-                        f"blocks_{current_stage_idx}_"
-                        f"{current_block_idx + current_layer_idx}"
-                    )
+                    common_kwargs = {
+                        "bn_epsilon": bn_epsilon,
+                        "padding": padding,
+                        "name": (
+                            f"blocks_{current_stage_idx}_"
+                            f"{current_block_idx + current_layer_idx}"
+                        ),
+                    }
                     if block_type == "ds":
                         x = apply_depthwise_separation_block(
-                            x, c, k, 1, s, se, act, name=name
+                            x, c, k, 1, s, se, act, **common_kwargs
                         )
                     elif block_type == "ir":
                         x = apply_inverted_residual_block(
-                            x, c, k, 1, 1, s, e, se, act, name=name
+                            x, c, k, 1, 1, s, e, se, act, **common_kwargs
                         )
                     elif block_type == "cn":
                         x = apply_conv2d_block(
@@ -282,7 +305,7 @@ def __init__(
                             kernel_size=k,
                             strides=s,
                             activation=act,
-                            name=name,
+                            **common_kwargs,
                         )
                     current_stride *= s
             features[f"BLOCK{current_stage_idx}_S{current_stride}"] = x
@@ -508,6 +531,7 @@ def __init__(
         name: str = "MobileNet100V3SmallMinimal",
         **kwargs,
     ):
+        # default to TF configuration (bn_epsilon=1e-3 and padding="same")
         super().__init__(
             1.0,
             1.0,
@@ -524,6 +548,8 @@ def __init__(
             config,
             minimal=True,
             name=name,
+            bn_epsilon=1e-3,
+            padding="same",
             **kwargs,
         )
 
@@ -598,6 +624,7 @@ def __init__(
         name: str = "MobileNet100V3LargeMinimal",
         **kwargs,
     ):
+        # default to TF configuration (bn_epsilon=1e-3 and padding="same")
         super().__init__(
             1.0,
             1.0,
@@ -614,6 +641,8 @@ def __init__(
             config,
             minimal=True,
             name=name,
+            bn_epsilon=1e-3,
+            padding="same",
             **kwargs,
         )
 
diff --git a/pyproject.toml b/pyproject.toml
index d8051ca..dff66b0 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -83,7 +83,7 @@ filterwarnings = [
 
 [tool.coverage.run]
 source = ["kimm"]
-omit = ["**/__init__.py", "*test*"]
+omit = ["**/__init__.py", "*test*", "kimm/utils/timm_utils.py"]
 
 [tool.coverage.report]
 exclude_lines = [
diff --git a/tools/convert_efficientnet_from_timm.py b/tools/convert_efficientnet_from_timm.py
new file mode 100644
index 0000000..1257e05
--- /dev/null
+++ b/tools/convert_efficientnet_from_timm.py
@@ -0,0 +1,183 @@
+"""
+pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu
+pip install timm
+"""
+import keras
+import numpy as np
+import timm
+import torch
+
+from kimm.models import efficientnet
+from kimm.utils.timm_utils import assign_weights
+from kimm.utils.timm_utils import is_same_weights
+from kimm.utils.timm_utils import separate_keras_weights
+from kimm.utils.timm_utils import separate_torch_state_dict
+
+timm_model_names = [
+    "tf_efficientnet_b0.ns_jft_in1k",
+    "tf_efficientnet_b1.ns_jft_in1k",
+    "tf_efficientnet_b2.ns_jft_in1k",
+    "tf_efficientnet_b3.ns_jft_in1k",
+    "tf_efficientnet_b4.ns_jft_in1k",
+    "tf_efficientnet_b5.ns_jft_in1k",
+    "tf_efficientnet_b6.ns_jft_in1k",
+    "tf_efficientnet_b7.ns_jft_in1k",
+    "tf_efficientnet_lite0.in1k",
+    "tf_efficientnet_lite1.in1k",
+    "tf_efficientnet_lite2.in1k",
+    "tf_efficientnet_lite3.in1k",
+    "tf_efficientnet_lite4.in1k",
+    "tf_efficientnetv2_s.in21k_ft_in1k",
+    "tf_efficientnetv2_m.in21k_ft_in1k",
+    "tf_efficientnetv2_l.in21k_ft_in1k",
+    "tf_efficientnetv2_xl.in21k_ft_in1k",
+    "tf_efficientnetv2_b0.in1k",
+    "tf_efficientnetv2_b1.in1k",
+    "tf_efficientnetv2_b2.in1k",
+    # "tf_efficientnetv2_b3.in1k",
+]
+keras_model_classes = [
+    efficientnet.EfficientNetB0,
+    efficientnet.EfficientNetB1,
+    efficientnet.EfficientNetB2,
+    efficientnet.EfficientNetB3,
+    efficientnet.EfficientNetB4,
+    efficientnet.EfficientNetB5,
+    efficientnet.EfficientNetB6,
+    efficientnet.EfficientNetB7,
+    efficientnet.EfficientNetLiteB0,
+    efficientnet.EfficientNetLiteB1,
+    efficientnet.EfficientNetLiteB2,
+    efficientnet.EfficientNetLiteB3,
+    efficientnet.EfficientNetLiteB4,
+    efficientnet.EfficientNetV2S,
+    efficientnet.EfficientNetV2M,
+    efficientnet.EfficientNetV2L,
+    efficientnet.EfficientNetV2XL,
+    efficientnet.EfficientNetV2B0,
+    efficientnet.EfficientNetV2B1,
+    efficientnet.EfficientNetV2B2,
+    # efficientnet.EfficientNetV2B3,
+]
+
+for timm_model_name, keras_model_class in zip(
+    timm_model_names, keras_model_classes
+):
+    """
+    Prepare timm model and keras model
+    """
+    input_shape = [224, 224, 3]
+    torch_model = timm.create_model(timm_model_name, pretrained=True)
+    torch_model = torch_model.eval()
+    trainable_state_dict, non_trainable_state_dict = separate_torch_state_dict(
+        torch_model.state_dict()
+    )
+    keras_model = keras_model_class(
+        input_shape=input_shape,
+        include_preprocessing=False,
+        classifier_activation="linear",
+    )
+    trainable_weights, non_trainable_weights = separate_keras_weights(
+        keras_model
+    )
+
+    # for torch_name, (_, keras_name) in zip(
+    #     trainable_state_dict.keys(), trainable_weights
+    # ):
+    #     print(f"{torch_name}    {keras_name}")
+
+    # print(len(trainable_state_dict.keys()))
+    # print(len(trainable_weights))
+
+    # exit()
+
+    """
+    Assign weights
+    """
+    for keras_weight, keras_name in trainable_weights + non_trainable_weights:
+        keras_name: str
+        torch_name = keras_name
+        torch_name = torch_name.replace("_", ".")
+        # stem
+        torch_name = torch_name.replace("conv.stem.conv2d", "conv_stem")
+        torch_name = torch_name.replace("conv.stem.bn", "bn1")
+        # blocks
+        if "EfficientNetV2" in keras_model_class.__name__:
+            if "blocks.0" in torch_name:
+                # normal conv
+                torch_name = torch_name.replace("conv2d", "conv")
+                torch_name = torch_name.replace("bn", "bn1")
+            elif "blocks.1" in torch_name or "blocks.2" in torch_name:
+                # edge residual block
+                torch_name = torch_name.replace("conv.exp.conv2d", "conv_exp")
+                torch_name = torch_name.replace("conv.exp.bn", "bn1")
+                torch_name = torch_name.replace("conv.pwl.conv2d", "conv_pwl")
+                torch_name = torch_name.replace("conv.pwl.bn", "bn2")
+        else:
+            if "blocks.0" in torch_name:
+                # depthwise separation block
+                torch_name = torch_name.replace("conv.dw.dwconv2d", "conv_dw")
+                torch_name = torch_name.replace("conv.dw.bn", "bn1")
+                torch_name = torch_name.replace("conv.pw.conv2d", "conv_pw")
+                torch_name = torch_name.replace("conv.pw.bn", "bn2")
+        # inverted residual block
+        torch_name = torch_name.replace("conv.pw.conv2d", "conv_pw")
+        torch_name = torch_name.replace("conv.pw.bn", "bn1")
+        torch_name = torch_name.replace("conv.dw.dwconv2d", "conv_dw")
+        torch_name = torch_name.replace("conv.dw.bn", "bn2")
+        torch_name = torch_name.replace("conv.pwl.conv2d", "conv_pwl")
+        torch_name = torch_name.replace("conv.pwl.bn", "bn3")
+        # se
+        torch_name = torch_name.replace("se.conv.reduce", "se.conv_reduce")
+        torch_name = torch_name.replace("se.conv.expand", "se.conv_expand")
+        # conv head
+        torch_name = torch_name.replace("conv.head.conv2d", "conv_head")
+        torch_name = torch_name.replace("conv.head.bn", "bn2")
+
+        # weights naming mapping
+        torch_name = torch_name.replace("kernel", "weight")  # conv2d
+        torch_name = torch_name.replace("gamma", "weight")  # bn
+        torch_name = torch_name.replace("beta", "bias")  # bn
+        torch_name = torch_name.replace("moving.mean", "running_mean")  # bn
+        torch_name = torch_name.replace("moving.variance", "running_var")  # bn
+
+        # assign weights
+        if torch_name in trainable_state_dict:
+            torch_weights = trainable_state_dict[torch_name].numpy()
+        elif torch_name in non_trainable_state_dict:
+            torch_weights = non_trainable_state_dict[torch_name].numpy()
+        else:
+            raise ValueError(
+                "Can't find the corresponding torch weights. "
+                f"Got keras_name={keras_name}, torch_name={torch_name}"
+            )
+        if is_same_weights(keras_name, keras_weight, torch_name, torch_weights):
+            assign_weights(keras_name, keras_weight, torch_weights)
+        else:
+            raise ValueError(
+                "Can't find the corresponding torch weights. The shape is "
+                f"mismatched. Got keras_name={keras_name}, "
+                f"keras_weight shape={keras_weight.shape}, "
+                f"torch_name={torch_name}, "
+                f"torch_weights shape={torch_weights.shape}"
+            )
+
+    """
+    Verify model outputs
+    """
+    np.random.seed(2023)
+    keras_data = np.random.uniform(size=[1] + input_shape).astype("float32")
+    torch_data = torch.from_numpy(np.transpose(keras_data, [0, 3, 1, 2]))
+    torch_y = torch_model(torch_data)
+    keras_y = keras_model(keras_data, training=False)
+    torch_y = torch_y.detach().cpu().numpy()
+    keras_y = keras.ops.convert_to_numpy(keras_y)
+    np.testing.assert_allclose(torch_y, keras_y, atol=2e-5)
+    print(f"{keras_model_class.__name__}: output matched!")
+
+    """
+    Save converted model
+    """
+    export_path = f"exported/{keras_model.name.lower()}_imagenet.keras"
+    keras_model.save(export_path)
+    print(f"Export to {export_path}")
diff --git a/tools/convert_mobilenet_v3_from_timm.py b/tools/convert_mobilenet_v3_from_timm.py
index 45063ae..bca1e62 100644
--- a/tools/convert_mobilenet_v3_from_timm.py
+++ b/tools/convert_mobilenet_v3_from_timm.py
@@ -16,18 +16,18 @@
 timm_model_names = [
     "mobilenetv3_small_050.lamb_in1k",
     "mobilenetv3_small_075.lamb_in1k",
-    # "tf_mobilenetv3_small_minimal_100.in1k",
+    "tf_mobilenetv3_small_minimal_100.in1k",
     "mobilenetv3_small_100.lamb_in1k",
     "mobilenetv3_large_100.miil_in21k_ft_in1k",
-    # "tf_mobilenetv3_large_minimal_100.in1k",
+    "tf_mobilenetv3_large_minimal_100.in1k",
 ]
 keras_model_classes = [
     mobilenet_v3.MobileNet050V3Small,
     mobilenet_v3.MobileNet075V3Small,
-    # mobilenet_v3.MobileNet100V3SmallMinimal,
+    mobilenet_v3.MobileNet100V3SmallMinimal,
     mobilenet_v3.MobileNet100V3Small,
     mobilenet_v3.MobileNet100V3Large,
-    # mobilenet_v3.MobileNet100V3LargeMinimal,
+    mobilenet_v3.MobileNet100V3LargeMinimal,
 ]
 
 for timm_model_name, keras_model_class in zip(


Test patch:
diff --git a/kimm/layers/attention_test.py b/kimm/layers/attention_test.py
index 39a37a3..639f238 100644
--- a/kimm/layers/attention_test.py
+++ b/kimm/layers/attention_test.py
@@ -1,15 +1,18 @@
 from absl.testing import parameterized
-from keras import random
 from keras.src import testing
 
 from kimm.layers.attention import Attention
 
 
 class AttentionTest(testing.TestCase, parameterized.TestCase):
-    def test_attention(self):
-        x = random.uniform([1, 197, 768])
-        layer = Attention(768)
-
-        y = layer(x)
-
-        self.assertEqual(y.shape, [1, 197, 768])
+    def test_attention_basic(self):
+        self.run_layer_test(
+            Attention,
+            init_kwargs={"hidden_dim": 20, "num_heads": 2},
+            input_shape=(1, 10, 20),
+            expected_output_shape=(1, 10, 20),
+            expected_num_trainable_weights=3,
+            expected_num_non_trainable_weights=0,
+            expected_num_losses=0,
+            supports_masking=False,
+        )
diff --git a/kimm/layers/layer_scale_test.py b/kimm/layers/layer_scale_test.py
index 5f7a46f..1a3fd89 100644
--- a/kimm/layers/layer_scale_test.py
+++ b/kimm/layers/layer_scale_test.py
@@ -1,15 +1,18 @@
 from absl.testing import parameterized
-from keras import random
 from keras.src import testing
 
 from kimm.layers.layer_scale import LayerScale
 
 
 class LayerScaleTest(testing.TestCase, parameterized.TestCase):
-    def test_layer_scale(self):
-        x = random.uniform([1, 123])
-        layer = LayerScale(123)
-
-        y = layer(x)
-
-        self.assertEqual(y.shape, [1, 123])
+    def test_layer_scale_basic(self):
+        self.run_layer_test(
+            LayerScale,
+            init_kwargs={"hidden_size": 10},
+            input_shape=(1, 10),
+            expected_output_shape=(1, 10),
+            expected_num_trainable_weights=1,
+            expected_num_non_trainable_weights=0,
+            expected_num_losses=0,
+            supports_masking=False,
+        )
diff --git a/kimm/layers/position_embedding_test.py b/kimm/layers/position_embedding_test.py
index d6f9363..9796e9f 100644
--- a/kimm/layers/position_embedding_test.py
+++ b/kimm/layers/position_embedding_test.py
@@ -1,15 +1,26 @@
 from absl.testing import parameterized
-from keras import random
+from keras import layers
 from keras.src import testing
 
 from kimm.layers.position_embedding import PositionEmbedding
 
 
 class PositionEmbeddingTest(testing.TestCase, parameterized.TestCase):
-    def test_position_embedding(self):
-        x = random.uniform([1, 123, 768])
-        layer = PositionEmbedding()
+    def test_position_embedding_basic(self):
+        self.run_layer_test(
+            PositionEmbedding,
+            init_kwargs={},
+            input_shape=(1, 10, 10),
+            expected_output_shape=(1, 11, 10),
+            expected_num_trainable_weights=2,
+            expected_num_non_trainable_weights=0,
+            expected_num_losses=0,
+            supports_masking=False,
+        )
 
-        y = layer(x)
-
-        self.assertEqual(y.shape, [1, 124, 768])
+    def test_position_embedding_invalid_input_shape(self):
+        inputs = layers.Input([3])
+        with self.assertRaisesRegex(
+            ValueError, "PositionEmbedding only accepts 3-dimensional input."
+        ):
+            PositionEmbedding()(inputs)
diff --git a/kimm/models/efficientnet_test.py b/kimm/models/efficientnet_test.py
new file mode 100644
index 0000000..8c240bc
--- /dev/null
+++ b/kimm/models/efficientnet_test.py
@@ -0,0 +1,114 @@
+from absl.testing import parameterized
+from keras import random
+from keras.src import testing
+
+from kimm.models.efficientnet import EfficientNetB0
+from kimm.models.efficientnet import EfficientNetB2
+from kimm.models.efficientnet import EfficientNetLiteB0
+from kimm.models.efficientnet import EfficientNetLiteB2
+from kimm.models.efficientnet import EfficientNetV2B0
+from kimm.models.efficientnet import EfficientNetV2S
+from kimm.utils import make_divisible
+
+
+class EfficientNetTest(testing.TestCase, parameterized.TestCase):
+    @parameterized.named_parameters(
+        [
+            (EfficientNetB0.__name__, EfficientNetB0, 224),
+            (EfficientNetB2.__name__, EfficientNetB2, 260),
+            (EfficientNetLiteB0.__name__, EfficientNetLiteB0, 224),
+            (EfficientNetLiteB2.__name__, EfficientNetLiteB2, 260),
+            (EfficientNetV2S.__name__, EfficientNetV2S, 300),
+            (EfficientNetV2B0.__name__, EfficientNetV2B0, 192),
+        ]
+    )
+    def test_efficentnet_base(self, model_class, image_size):
+        # TODO: test the correctness of the real image
+        x = random.uniform([1, image_size, image_size, 3]) * 255.0
+        model = model_class()
+
+        y = model.predict(x)
+
+        self.assertEqual(y.shape, (1, 1000))
+
+    @parameterized.named_parameters(
+        [
+            (EfficientNetB0.__name__, EfficientNetB0, 1.0),
+            (EfficientNetB2.__name__, EfficientNetB2, 1.1),
+            (EfficientNetLiteB0.__name__, EfficientNetLiteB0, 1.0),
+            (EfficientNetLiteB2.__name__, EfficientNetLiteB2, 1.1),
+        ]
+    )
+    def test_efficentnet_feature_extractor(self, model_class, width):
+        x = random.uniform([1, 224, 224, 3]) * 255.0
+        model = model_class(
+            input_shape=[224, 224, 3], as_feature_extractor=True
+        )
+
+        y = model.predict(x)
+
+        self.assertIsInstance(y, dict)
+        self.assertAllEqual(
+            list(y.keys()), model_class.available_feature_keys()
+        )
+        self.assertEqual(
+            list(y["STEM_S2"].shape), [1, 112, 112, make_divisible(32 * width)]
+        )
+        self.assertEqual(
+            list(y["BLOCK1_S4"].shape), [1, 56, 56, make_divisible(24 * width)]
+        )
+        self.assertEqual(
+            list(y["BLOCK2_S8"].shape), [1, 28, 28, make_divisible(40 * width)]
+        )
+        self.assertEqual(
+            list(y["BLOCK3_S16"].shape), [1, 14, 14, make_divisible(80 * width)]
+        )
+        self.assertEqual(
+            list(y["BLOCK5_S32"].shape), [1, 7, 7, make_divisible(192 * width)]
+        )
+
+    @parameterized.named_parameters(
+        [
+            (EfficientNetV2S.__name__, EfficientNetV2S, 1.0),
+            (EfficientNetV2B0.__name__, EfficientNetV2B0, 1.0),
+        ]
+    )
+    def test_efficentnet_v2_feature_extractor(self, model_class, width):
+        x = random.uniform([1, 224, 224, 3]) * 255.0
+        model = model_class(
+            input_shape=[224, 224, 3], as_feature_extractor=True
+        )
+
+        y = model.predict(x)
+
+        self.assertIsInstance(y, dict)
+        self.assertAllEqual(
+            list(y.keys()), model_class.available_feature_keys()
+        )
+        if "EfficientNetV2S" in model_class.__name__:
+            self.assertEqual(list(y["STEM_S2"].shape), [1, 112, 112, 24])
+            self.assertEqual(list(y["BLOCK1_S4"].shape), [1, 56, 56, 48])
+            self.assertEqual(list(y["BLOCK2_S8"].shape), [1, 28, 28, 64])
+            self.assertEqual(list(y["BLOCK3_S16"].shape), [1, 14, 14, 128])
+            self.assertEqual(list(y["BLOCK5_S32"].shape), [1, 7, 7, 256])
+        elif "EfficientNetV2B" in model_class.__name__:
+            self.assertEqual(
+                list(y["STEM_S2"].shape),
+                [1, 112, 112, make_divisible(32 * width)],
+            )
+            self.assertEqual(
+                list(y["BLOCK1_S4"].shape),
+                [1, 56, 56, make_divisible(32 * width)],
+            )
+            self.assertEqual(
+                list(y["BLOCK2_S8"].shape),
+                [1, 28, 28, make_divisible(48 * width)],
+            )
+            self.assertEqual(
+                list(y["BLOCK3_S16"].shape),
+                [1, 14, 14, make_divisible(96 * width)],
+            )
+            self.assertEqual(
+                list(y["BLOCK5_S32"].shape),
+                [1, 7, 7, make_divisible(192 * width)],
+            )
diff --git a/kimm/models/mobilenet_v3_test.py b/kimm/models/mobilenet_v3_test.py
index 9d40fa8..63281f6 100644
--- a/kimm/models/mobilenet_v3_test.py
+++ b/kimm/models/mobilenet_v3_test.py
@@ -7,7 +7,7 @@
 from kimm.utils import make_divisible
 
 
-class MobileNetV2Test(testing.TestCase, parameterized.TestCase):
+class MobileNetV3Test(testing.TestCase, parameterized.TestCase):
     @parameterized.named_parameters(
         [
             (MobileNet100V3Small.__name__, MobileNet100V3Small),