Repository: keras-team/keras-cv

Issue:
Make MobileNetV3 a Functional subclass model & remove closures
This should follow the example of #1412 and #1401, including

- Docstring updates
- Colab demonstrating that weight loading works (if pre-trained weights exist) and numerics are identical

Discussion:
ID6109: The changes have been made. I'll implement the same in my other PRs. Do let me know if there are further corrections needed. Thanks.

Patch:
diff --git a/keras_cv/models/efficientnet_v1.py b/keras_cv/models/efficientnet_v1.py
index def28f8..c66d1d3 100644
--- a/keras_cv/models/efficientnet_v1.py
+++ b/keras_cv/models/efficientnet_v1.py
@@ -25,8 +25,9 @@ import copy
 import math
 
 import tensorflow as tf
-from keras import backend
-from keras import layers
+from tensorflow import keras
+from tensorflow.keras import backend
+from tensorflow.keras import layers
 
 from keras_cv.models import utils
 from keras_cv.models.weights import parse_weights
@@ -128,7 +129,7 @@ BASE_DOCSTRING = """Instantiates the {name} architecture.
     - [EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks](
         https://arxiv.org/abs/1905.11946) (ICML 2019)
 
-    This function returns a Keras image classification model.
+    This class represents a Keras image classification model.
 
     For image classification use cases, see
     [this page for detailed examples](
@@ -139,13 +140,13 @@ BASE_DOCSTRING = """Instantiates the {name} architecture.
     https://keras.io/guides/transfer_learning/).
 
     Args:
-        include_rescaling: whether or not to Rescale the inputs.If set to True,
+        include_rescaling: bool, whether or not to Rescale the inputs.If set to True,
                     inputs will be passed through a `Rescaling(1/255.0)` layer.
-        include_top: Whether to include the fully-connected
+        include_top: bool, Whether to include the fully-connected
             layer at the top of the network.
         weights: One of `None` (random initialization),
                 or the path to the weights file to be loaded.
-        input_shape: Optional shape tuple.
+        input_shape: tuple, Optional shape tuple.
             It should have exactly 3 inputs channels.
         input_tensor: optional Keras tensor (i.e. output of `layers.Input()`)
             to use as image input for the model.
@@ -160,7 +161,7 @@ BASE_DOCSTRING = """Instantiates the {name} architecture.
                 the output of the model will be a 2D tensor.
             - `max` means that global max pooling will
                 be applied.
-        num_classes: Optional number of classes to classify images
+        num_classes: int, Optional number of classes to classify images
             into, only to be specified if `include_top` is True, and
             if no `weights` argument is specified. Defaults to None.
         classifier_activation: A `str` or callable. The activation function to use
@@ -174,7 +175,6 @@ BASE_DOCSTRING = """Instantiates the {name} architecture.
         A `keras.Model` instance.
 """
 
-
 BN_AXIS = 3
 
 
@@ -201,151 +201,217 @@ def correct_pad(inputs, kernel_size):
     )
 
 
-def EfficientNetBlock(
+def apply_conv_bn(
+    x,
+    conv_type,
+    filters,
+    kernel_size,
+    strides=1,
+    padding="same",
+    use_bias=False,
+    kernel_initializer=CONV_KERNEL_INITIALIZER,
+    bn_norm=True,
     activation="swish",
-    drop_rate=0.0,
     name="",
+):
+    """
+    Represents Convolutional Block with optional Batch Normalization layer and activation layer
+
+    Args:
+        x: Tensor
+        conv_type: str, Type of Conv layer to be used in block.
+            - 'normal': The Conv2D layer will be used.
+            - 'depth': The DepthWiseConv2D layer will be used.
+        filters: int, The filter size of the Conv layer.
+                It should be `None` when `conv_type` is set as `depth`
+        kernel_size: int (or) tuple, The kernel size of the Conv layer.
+        strides: int (or) tuple, The stride value of Conv layer.
+        padding: str (or) callable, The type of padding for Conv layer.
+        use_bias: bool, Boolean to use bias for Conv layer.
+        kernel_initializer: dict (or) str (or) callable, The kernel initializer for Conv layer.
+        bn_norm: bool, Boolean to add BatchNormalization layer after Conv layer.
+        activation: str (or) callable, Activation to be applied on the output at end.
+        name: name of the block
+
+    Returns:
+        tf.Tensor
+    """
+    if conv_type == "normal":
+        if filters is None or kernel_size is None:
+            raise ValueError(
+                "The filter size and kernel size should be set for Conv2D layer."
+            )
+        x = layers.Conv2D(
+            filters,
+            kernel_size,
+            strides=strides,
+            padding=padding,
+            use_bias=use_bias,
+            kernel_initializer=kernel_initializer,
+            name=name + "_conv",
+        )(x)
+    elif conv_type == "depth":
+        if filters is not None:
+            raise ValueError(
+                "Filter size shouldn't be set for DepthWiseConv2D layer."
+            )
+        if kernel_size is None or strides is None:
+            raise ValueError(
+                "The kernel size and strides should be set for DepthWiseConv2D layer."
+            )
+        x = layers.DepthwiseConv2D(
+            kernel_size,
+            strides=strides,
+            padding=padding,
+            use_bias=use_bias,
+            depthwise_initializer=kernel_initializer,
+            name=name + "_dwconv",
+        )(x)
+    else:
+        raise ValueError(
+            "The 'conv_type' parameter should be set either to 'normal' or 'depth'"
+        )
+
+    if bn_norm:
+        x = layers.BatchNormalization(axis=BN_AXIS, name=name + "_bn")(x)
+    if activation is not None:
+        x = layers.Activation(activation, name=name + "_activation")(x)
+
+    return x
+
+
+def apply_efficientnet_block(
+    inputs,
     filters_in=32,
     filters_out=16,
     kernel_size=3,
     strides=1,
+    activation="swish",
     expand_ratio=1,
     se_ratio=0.0,
     id_skip=True,
+    drop_rate=0.0,
+    name="",
 ):
     """An inverted residual block.
 
     Args:
-        inputs: input tensor.
-        activation: activation function.
-        drop_rate: float between 0 and 1, fraction of the input units to drop.
-        name: string, block label.
+        inputs: Tensor, The input tensor of the block
         filters_in: integer, the number of input filters.
         filters_out: integer, the number of output filters.
         kernel_size: integer, the dimension of the convolution window.
         strides: integer, the stride of the convolution.
+        activation: activation function.
         expand_ratio: integer, scaling coefficient for the input filters.
         se_ratio: float between 0 and 1, fraction to squeeze the input filters.
         id_skip: boolean.
+        drop_rate: float between 0 and 1, fraction of the input units to drop.
+        name: string, block label.
 
     Returns:
-        output tensor for the block.
+        tf.Tensor
     """
+    filters = filters_in * expand_ratio
+    if expand_ratio != 1:
+        x = apply_conv_bn(
+            x=inputs,
+            conv_type="normal",
+            filters=filters,
+            kernel_size=1,
+            padding="same",
+            use_bias=False,
+            kernel_initializer=CONV_KERNEL_INITIALIZER,
+            bn_norm=True,
+            activation=activation,
+            name=name + "_expand",
+        )
+    else:
+        x = inputs
 
-    # Expansion phase
-    def apply(inputs):
-        filters = filters_in * expand_ratio
-        if expand_ratio != 1:
-            x = layers.Conv2D(
-                filters,
-                1,
-                padding="same",
-                use_bias=False,
-                kernel_initializer=CONV_KERNEL_INITIALIZER,
-                name=name + "expand_conv",
-            )(inputs)
-            x = layers.BatchNormalization(
-                axis=BN_AXIS, name=name + "expand_bn"
-            )(x)
-            x = layers.Activation(activation, name=name + "expand_activation")(
-                x
-            )
-        else:
-            x = inputs
+    # Depthwise Convolution
+    if strides == 2:
+        x = layers.ZeroPadding2D(
+            padding=correct_pad(x, kernel_size),
+            name=name + "_dwconv_pad",
+        )(x)
+        conv_pad = "valid"
+    else:
+        conv_pad = "same"
+
+    x = apply_conv_bn(
+        x=x,
+        conv_type="depth",
+        filters=None,
+        kernel_size=kernel_size,
+        strides=strides,
+        padding=conv_pad,
+        use_bias=False,
+        kernel_initializer=CONV_KERNEL_INITIALIZER,
+        bn_norm=True,
+        activation=activation,
+        name=name,
+    )
 
-        # Depthwise Convolution
-        if strides == 2:
-            x = layers.ZeroPadding2D(
-                padding=correct_pad(x, kernel_size),
-                name=name + "dwconv_pad",
-            )(x)
-            conv_pad = "valid"
+    # Squeeze and Excitation phase
+    if 0 < se_ratio <= 1:
+        filters_se = max(1, int(filters_in * se_ratio))
+        se = layers.GlobalAveragePooling2D(name=name + "_se_squeeze")(x)
+        if BN_AXIS == 1:
+            se_shape = (filters, 1, 1)
         else:
-            conv_pad = "same"
-        x = layers.DepthwiseConv2D(
-            kernel_size,
-            strides=strides,
-            padding=conv_pad,
-            use_bias=False,
-            depthwise_initializer=CONV_KERNEL_INITIALIZER,
-            name=name + "dwconv",
-        )(x)
-        x = layers.BatchNormalization(axis=BN_AXIS, name=name + "bn")(x)
-        x = layers.Activation(activation, name=name + "activation")(x)
-
-        # Squeeze and Excitation phase
-        if 0 < se_ratio <= 1:
-            filters_se = max(1, int(filters_in * se_ratio))
-            se = layers.GlobalAveragePooling2D(name=name + "se_squeeze")(x)
-            if BN_AXIS == 1:
-                se_shape = (filters, 1, 1)
-            else:
-                se_shape = (1, 1, filters)
-            se = layers.Reshape(se_shape, name=name + "se_reshape")(se)
-            se = layers.Conv2D(
-                filters_se,
-                1,
-                padding="same",
-                activation=activation,
-                kernel_initializer=CONV_KERNEL_INITIALIZER,
-                name=name + "se_reduce",
-            )(se)
-            se = layers.Conv2D(
-                filters,
-                1,
-                padding="same",
-                activation="sigmoid",
-                kernel_initializer=CONV_KERNEL_INITIALIZER,
-                name=name + "se_expand",
-            )(se)
-            x = layers.multiply([x, se], name=name + "se_excite")
-
-        # Output phase
-        x = layers.Conv2D(
-            filters_out,
+            se_shape = (1, 1, filters)
+        se = layers.Reshape(se_shape, name=name + "_se_reshape")(se)
+        se = layers.Conv2D(
+            filters_se,
             1,
             padding="same",
-            use_bias=False,
+            activation=activation,
             kernel_initializer=CONV_KERNEL_INITIALIZER,
-            name=name + "project_conv",
-        )(x)
-        x = layers.BatchNormalization(axis=BN_AXIS, name=name + "project_bn")(x)
-        if id_skip and strides == 1 and filters_in == filters_out:
-            if drop_rate > 0:
-                x = layers.Dropout(
-                    drop_rate, noise_shape=(None, 1, 1, 1), name=name + "drop"
-                )(x)
-            x = layers.add([x, inputs], name=name + "add")
-        return x
+            name=name + "_se_reduce",
+        )(se)
+        se = layers.Conv2D(
+            filters,
+            1,
+            padding="same",
+            activation="sigmoid",
+            kernel_initializer=CONV_KERNEL_INITIALIZER,
+            name=name + "_se_expand",
+        )(se)
+        x = layers.multiply([x, se], name=name + "_se_excite")
+
+    # Output phase
+    x = apply_conv_bn(
+        x=x,
+        conv_type="normal",
+        filters=filters_out,
+        kernel_size=1,
+        padding="same",
+        use_bias=False,
+        kernel_initializer=CONV_KERNEL_INITIALIZER,
+        bn_norm=True,
+        activation=None,
+        name=name + "_project",
+    )
 
-    return apply
+    if id_skip and strides == 1 and filters_in == filters_out:
+        if drop_rate > 0:
+            x = layers.Dropout(
+                drop_rate,
+                noise_shape=(None, 1, 1, 1),
+                name=name + "_drop",
+            )(x)
+        x = layers.add([x, inputs], name=name + "_add")
 
+    return x
 
-def EfficientNet(
-    include_rescaling,
-    include_top,
-    width_coefficient,
-    depth_coefficient,
-    default_size,
-    dropout_rate=0.2,
-    drop_connect_rate=0.2,
-    depth_divisor=8,
-    activation="swish",
-    blocks_args="default",
-    model_name="efficientnet",
-    weights=None,
-    input_shape=(None, None, 3),
-    input_tensor=None,
-    pooling=None,
-    num_classes=None,
-    classifier_activation="softmax",
-):
-    """Instantiates the EfficientNet architecture using given scaling coefficients.
 
+@keras.utils.register_keras_serializable(package="keras_cv.models")
+class EfficientNet(keras.Model):
+    """This class represents a Keras EfficientNet architecture.
     Args:
-        include_rescaling: whether or not to Rescale the inputs.If set to True,
+        include_rescaling: bool, whether or not to Rescale the inputs.If set to True,
                 inputs will be passed through a `Rescaling(1/255.0)` layer.
-        include_top: whether to include the fully-connected
+        include_top: bool, whether to include the fully-connected
             layer at the top of the network.
         width_coefficient: float, scaling coefficient for network width.
         depth_coefficient: float, scaling coefficient for network depth.
@@ -379,41 +445,192 @@ def EfficientNet(
         classifier_activation: A `str` or callable. The activation function to use
             on the "top" layer. Ignored unless `include_top=True`. Set
             `classifier_activation=None` to return the logits of the "top" layer.
-
     Returns:
       A `keras.Model` instance.
-
     Raises:
       ValueError: in case of invalid argument for `weights`,
         or invalid input shape.
       ValueError: if `classifier_activation` is not `softmax` or `None` when
         using a pretrained top layer.
     """
-    if blocks_args == "default":
-        blocks_args = DEFAULT_BLOCKS_ARGS
 
-    if weights and not tf.io.gfile.exists(weights):
-        raise ValueError(
-            "The `weights` argument should be either `None` or the path to the "
-            "weights file to be loaded. Weights file not found at location: {weights}"
-        )
+    def __init__(
+        self,
+        include_rescaling,
+        include_top,
+        width_coefficient,
+        depth_coefficient,
+        default_size,
+        dropout_rate=0.2,
+        drop_connect_rate=0.2,
+        depth_divisor=8,
+        activation="swish",
+        blocks_args="default",
+        model_name="efficientnet",
+        weights=None,
+        input_shape=(None, None, 3),
+        input_tensor=None,
+        pooling=None,
+        num_classes=None,
+        classifier_activation="softmax",
+        **kwargs,
+    ):
+        blocks_args_type = blocks_args
 
-    if include_top and not num_classes:
-        raise ValueError(
-            "If `include_top` is True, you should specify `num_classes`. "
-            f"Received: num_classes={num_classes}"
+        if blocks_args == "default":
+            blocks_args = DEFAULT_BLOCKS_ARGS
+
+        if weights and not tf.io.gfile.exists(weights):
+            raise ValueError(
+                "The `weights` argument should be either `None` or the path to the "
+                "weights file to be loaded. Weights file not found at location: {weights}"
+            )
+
+        if include_top and not num_classes:
+            raise ValueError(
+                "If `include_top` is True, you should specify `num_classes`. "
+                f"Received: num_classes={num_classes}"
+            )
+
+        if include_top and pooling:
+            raise ValueError(
+                f"`pooling` must be `None` when `include_top=True`."
+                f"Received pooling={pooling} and include_top={include_top}. "
+            )
+
+        img_input = utils.parse_model_inputs(input_shape, input_tensor)
+
+        # Build stem
+        x = img_input
+
+        if include_rescaling:
+            # Use common rescaling strategy across keras_cv
+            x = layers.Rescaling(1.0 / 255.0)(x)
+
+        x = layers.ZeroPadding2D(
+            padding=correct_pad(x, 3), name="stem_conv_pad"
+        )(x)
+
+        x = apply_conv_bn(
+            x=x,
+            conv_type="normal",
+            filters=EfficientNet.round_filters(
+                32, width_coefficient, depth_divisor
+            ),
+            kernel_size=3,
+            strides=2,
+            padding="valid",
+            use_bias=False,
+            kernel_initializer=CONV_KERNEL_INITIALIZER,
+            bn_norm=True,
+            activation=activation,
+            name="stem",
         )
 
-    if include_top and pooling:
-        raise ValueError(
-            f"`pooling` must be `None` when `include_top=True`."
-            f"Received pooling={pooling} and include_top={include_top}. "
+        # Build blocks
+        blocks_args = copy.deepcopy(blocks_args)
+
+        b = 0
+        blocks = float(
+            sum(
+                EfficientNet.round_repeats(args["repeats"], depth_coefficient)
+                for args in blocks_args
+            )
         )
+        for i, args in enumerate(blocks_args):
+            assert args["repeats"] > 0
+            # Update block input and output filters based on depth multiplier.
+            args["filters_in"] = EfficientNet.round_filters(
+                args["filters_in"], width_coefficient, depth_divisor
+            )
+            args["filters_out"] = EfficientNet.round_filters(
+                args["filters_out"], width_coefficient, depth_divisor
+            )
 
-    img_input = utils.parse_model_inputs(input_shape, input_tensor)
+            for j in range(
+                EfficientNet.round_repeats(
+                    args.pop("repeats"), depth_coefficient
+                )
+            ):
+                # The first block needs to take care of stride and filter size
+                # increase.
+                if j > 0:
+                    args["strides"] = 1
+                    args["filters_in"] = args["filters_out"]
+                x = apply_efficientnet_block(
+                    inputs=x,
+                    activation=activation,
+                    drop_rate=drop_connect_rate * b / blocks,
+                    name="block{}{}".format(i + 1, chr(j + 97)),
+                    **args,
+                )
+                b += 1
+
+        # Build top
+        x = apply_conv_bn(
+            x=x,
+            conv_type="normal",
+            filters=self.round_filters(1280, width_coefficient, depth_divisor),
+            kernel_size=1,
+            padding="same",
+            use_bias=False,
+            kernel_initializer=CONV_KERNEL_INITIALIZER,
+            bn_norm=True,
+            activation=activation,
+            name="top",
+        )
 
-    def round_filters(filters, divisor=depth_divisor):
-        """Round number of filters based on depth multiplier."""
+        if include_top:
+            x = layers.GlobalAveragePooling2D(name="avg_pool")(x)
+            if dropout_rate > 0:
+                x = layers.Dropout(dropout_rate, name="top_dropout")(x)
+            x = layers.Dense(
+                num_classes,
+                activation=classifier_activation,
+                kernel_initializer=DENSE_KERNEL_INITIALIZER,
+                name="predictions",
+            )(x)
+        else:
+            if pooling == "avg":
+                x = layers.GlobalAveragePooling2D(name="avg_pool")(x)
+            elif pooling == "max":
+                x = layers.GlobalMaxPooling2D(name="max_pool")(x)
+
+        inputs = img_input
+
+        # Create model.
+        super().__init__(inputs=inputs, outputs=x, name=model_name, **kwargs)
+
+        # Load weights.
+        if weights is not None:
+            self.load_weights(weights)
+
+        self.include_rescaling = include_rescaling
+        self.include_top = include_top
+        self.width_coefficient = width_coefficient
+        self.depth_coefficient = depth_coefficient
+        self.default_size = default_size
+        self.dropout_rate = dropout_rate
+        self.drop_connect_rate = drop_connect_rate
+        self.depth_divisor = depth_divisor
+        self.activation = activation
+        self.blocks_args = blocks_args_type
+        self.input_tensor = input_tensor
+        self.pooling = pooling
+        self.num_classes = num_classes
+        self.classifier_activation = classifier_activation
+
+    @staticmethod
+    def round_filters(filters, width_coefficient, divisor):
+        """Round number of filters based on depth multiplier.
+        Args:
+            filters: int, number of filters for Conv layer
+            width_coefficient: float, denotes the scaling coefficient of network width
+            divisor: int, a unit of network width
+
+        Returns:
+            int, new rounded filters value for Conv layer
+        """
         filters *= width_coefficient
         new_filters = max(
             divisor, int(filters + divisor / 2) // divisor * divisor
@@ -423,92 +640,42 @@ def EfficientNet(
             new_filters += divisor
         return int(new_filters)
 
-    def round_repeats(repeats):
-        """Round number of repeats based on depth multiplier."""
-        return int(math.ceil(depth_coefficient * repeats))
-
-    # Build stem
-    x = img_input
-
-    if include_rescaling:
-        # Use common rescaling strategy across keras_cv
-        x = layers.Rescaling(1.0 / 255.0)(x)
-
-    x = layers.ZeroPadding2D(padding=correct_pad(x, 3), name="stem_conv_pad")(x)
-    x = layers.Conv2D(
-        round_filters(32),
-        3,
-        strides=2,
-        padding="valid",
-        use_bias=False,
-        kernel_initializer=CONV_KERNEL_INITIALIZER,
-        name="stem_conv",
-    )(x)
-    x = layers.BatchNormalization(axis=BN_AXIS, name="stem_bn")(x)
-    x = layers.Activation(activation, name="stem_activation")(x)
-
-    # Build blocks
-    blocks_args = copy.deepcopy(blocks_args)
-
-    b = 0
-    blocks = float(sum(round_repeats(args["repeats"]) for args in blocks_args))
-    for i, args in enumerate(blocks_args):
-        assert args["repeats"] > 0
-        # Update block input and output filters based on depth multiplier.
-        args["filters_in"] = round_filters(args["filters_in"])
-        args["filters_out"] = round_filters(args["filters_out"])
-
-        for j in range(round_repeats(args.pop("repeats"))):
-            # The first block needs to take care of stride and filter size
-            # increase.
-            if j > 0:
-                args["strides"] = 1
-                args["filters_in"] = args["filters_out"]
-            x = EfficientNetBlock(
-                activation,
-                drop_connect_rate * b / blocks,
-                name="block{}{}_".format(i + 1, chr(j + 97)),
-                **args,
-            )(x)
-            b += 1
-
-    # Build top
-    x = layers.Conv2D(
-        round_filters(1280),
-        1,
-        padding="same",
-        use_bias=False,
-        kernel_initializer=CONV_KERNEL_INITIALIZER,
-        name="top_conv",
-    )(x)
-    x = layers.BatchNormalization(axis=BN_AXIS, name="top_bn")(x)
-    x = layers.Activation(activation, name="top_activation")(x)
-    if include_top:
-        x = layers.GlobalAveragePooling2D(name="avg_pool")(x)
-        if dropout_rate > 0:
-            x = layers.Dropout(dropout_rate, name="top_dropout")(x)
-        x = layers.Dense(
-            num_classes,
-            activation=classifier_activation,
-            kernel_initializer=DENSE_KERNEL_INITIALIZER,
-            name="predictions",
-        )(x)
-    else:
-        if pooling == "avg":
-            x = layers.GlobalAveragePooling2D(name="avg_pool")(x)
-        elif pooling == "max":
-            x = layers.GlobalMaxPooling2D(name="max_pool")(x)
-
-    inputs = img_input
+    @staticmethod
+    def round_repeats(repeats, depth_coefficient):
+        """Round number of repeats based on depth multiplier.
+        Args:
+            repeats: int, number of repeats of efficentnet block
+            depth_coefficient: float, denotes the scaling coefficient of  network depth
 
-    # Create model.
-    model = tf.keras.Model(inputs, x, name=model_name)
-
-    # Load weights.
-    if weights is not None:
-        model.load_weights(weights)
+        Returns:
+            int, rounded repeats
+        """
+        return int(math.ceil(depth_coefficient * repeats))
 
-    return model
+    def get_config(self):
+        return {
+            "include_rescaling": self.include_rescaling,
+            "include_top": self.include_top,
+            "width_coefficient": self.width_coefficient,
+            "depth_coefficient": self.depth_coefficient,
+            "default_size": self.default_size,
+            "dropout_rate": self.dropout_rate,
+            "drop_connect_rate": self.drop_connect_rate,
+            "depth_divisor": self.depth_divisor,
+            "activation": self.activation,
+            "blocks_args": self.blocks_args,
+            "input_tensor": self.input_tensor,
+            "input_shape": self.input_shape[1:],
+            "model_name": self.name,
+            "pooling": self.pooling,
+            "num_classes": self.num_classes,
+            "classifier_activation": self.classifier_activation,
+            "trainable": self.trainable,
+        }
+
+    @classmethod
+    def from_config(cls, config):
+        return cls(**config)
 
 
 def EfficientNetB0(
@@ -521,6 +688,7 @@ def EfficientNetB0(
     input_tensor=None,
     pooling=None,
     classifier_activation="softmax",
+    name="efficientnetb0",
     **kwargs,
 ):
     return EfficientNet(
@@ -530,7 +698,7 @@ def EfficientNetB0(
         depth_coefficient=1.0,
         default_size=224,
         dropout_rate=0.2,
-        model_name="efficientnetb0",
+        model_name=name,
         weights=parse_weights(weights, include_top, "efficientnetb0"),
         input_shape=input_shape,
         input_tensor=input_tensor,
@@ -551,6 +719,7 @@ def EfficientNetB1(
     input_tensor=None,
     pooling=None,
     classifier_activation="softmax",
+    name="efficientnetb1",
     **kwargs,
 ):
     return EfficientNet(
@@ -560,7 +729,7 @@ def EfficientNetB1(
         depth_coefficient=1.1,
         default_size=240,
         dropout_rate=0.2,
-        model_name="efficientnetb1",
+        model_name=name,
         weights=parse_weights(weights, include_top, "efficientnetb1"),
         input_shape=input_shape,
         input_tensor=input_tensor,
@@ -581,6 +750,7 @@ def EfficientNetB2(
     input_tensor=None,
     pooling=None,
     classifier_activation="softmax",
+    name="efficientnetb2",
     **kwargs,
 ):
     return EfficientNet(
@@ -590,7 +760,7 @@ def EfficientNetB2(
         depth_coefficient=1.2,
         default_size=260,
         dropout_rate=0.3,
-        model_name="efficientnetb2",
+        model_name=name,
         weights=parse_weights(weights, include_top, "efficientnetb2"),
         input_shape=input_shape,
         input_tensor=input_tensor,
@@ -611,6 +781,7 @@ def EfficientNetB3(
     input_tensor=None,
     pooling=None,
     classifier_activation="softmax",
+    name="efficientnetb3",
     **kwargs,
 ):
     return EfficientNet(
@@ -620,7 +791,7 @@ def EfficientNetB3(
         depth_coefficient=1.4,
         default_size=300,
         dropout_rate=0.3,
-        model_name="efficientnetb3",
+        model_name=name,
         weights=parse_weights(weights, include_top, "efficientnetb3"),
         input_shape=input_shape,
         input_tensor=input_tensor,
@@ -641,6 +812,7 @@ def EfficientNetB4(
     input_tensor=None,
     pooling=None,
     classifier_activation="softmax",
+    name="efficientnetb4",
     **kwargs,
 ):
     return EfficientNet(
@@ -650,7 +822,7 @@ def EfficientNetB4(
         depth_coefficient=1.8,
         default_size=380,
         dropout_rate=0.4,
-        model_name="efficientnetb4",
+        model_name=name,
         weights=parse_weights(weights, include_top, "efficientnetb4"),
         input_shape=input_shape,
         input_tensor=input_tensor,
@@ -671,6 +843,7 @@ def EfficientNetB5(
     input_tensor=None,
     pooling=None,
     classifier_activation="softmax",
+    name="efficientnetb5",
     **kwargs,
 ):
     return EfficientNet(
@@ -680,7 +853,7 @@ def EfficientNetB5(
         depth_coefficient=2.2,
         default_size=456,
         dropout_rate=0.4,
-        model_name="efficientnetb5",
+        model_name=name,
         weights=parse_weights(weights, include_top, "efficientnetb5"),
         input_shape=input_shape,
         input_tensor=input_tensor,
@@ -701,6 +874,7 @@ def EfficientNetB6(
     input_tensor=None,
     pooling=None,
     classifier_activation="softmax",
+    name="efficientnetb6",
     **kwargs,
 ):
     return EfficientNet(
@@ -710,7 +884,7 @@ def EfficientNetB6(
         depth_coefficient=2.6,
         default_size=528,
         dropout_rate=0.5,
-        model_name="efficientnetb6",
+        model_name=name,
         weights=parse_weights(weights, include_top, "efficientnetb6"),
         input_shape=input_shape,
         input_tensor=input_tensor,
@@ -731,6 +905,7 @@ def EfficientNetB7(
     input_tensor=None,
     pooling=None,
     classifier_activation="softmax",
+    name="efficientnetb7",
     **kwargs,
 ):
     return EfficientNet(
@@ -740,7 +915,7 @@ def EfficientNetB7(
         depth_coefficient=3.1,
         default_size=600,
         dropout_rate=0.5,
-        model_name="efficientnetb7",
+        model_name=name,
         weights=parse_weights(weights, include_top, "efficientnetb7"),
         input_shape=input_shape,
         input_tensor=input_tensor,
diff --git a/keras_cv/models/mobilenet_v3.py b/keras_cv/models/mobilenet_v3.py
index 83e0c07..da0e602 100644
--- a/keras_cv/models/mobilenet_v3.py
+++ b/keras_cv/models/mobilenet_v3.py
@@ -36,24 +36,24 @@ BASE_DOCSTRING = """Instantiates the {name} architecture.
         - [Searching for MobileNetV3](https://arxiv.org/abs/1905.02244)
         - [Based on the Original keras.applications MobileNetv3](https://github.com/keras-team/keras/blob/master/keras/applications/mobilenet_v3.py)
 
-    This function returns a Keras {name} model.
+    This class represents a Keras {name} model.
 
     For transfer learning use cases, make sure to read the [guide to transfer
         learning & fine-tuning](https://keras.io/guides/transfer_learning/).
 
     Args:
-        include_rescaling: whether or not to Rescale the inputs.If set to True,
+        include_rescaling: bool, whether or not to rescale the inputs. If set to True,
             inputs will be passed through a `Rescaling(scale=1 / 255)`
-            layer, defaults to True.
-        include_top: whether to include the fully-connected layer at the top of the
-            network.  If provided, `num_classes` must be provided.
-        num_classes: optional number of classes to classify images into, only to be
+            layer. Defaults to True.
+        include_top: bool, whether to include the fully-connected layer at the top of the
+            network. If provided, `num_classes` must be provided.
+        num_classes: integer, optional number of classes to classify images into. Only to be
             specified if `include_top` is True, and if no `weights` argument is
             specified.
-        weights: one of `None` (random initialization), or a pretrained weight file
+        weights: one of `None` (random initialization) or a pretrained weight file
             path.
         input_shape: optional shape tuple, defaults to (None, None, 3).
-        input_tensor: optional Keras tensor (i.e. output of `layers.Input()`)
+        input_tensor: optional Keras tensor (i.e., output of `layers.Input()`)
             to use as image input for the model.
         pooling: optional pooling mode for feature extraction
             when `include_top` is `False`.
@@ -63,7 +63,7 @@ BASE_DOCSTRING = """Instantiates the {name} architecture.
                 of the last convolutional block, and thus the output of the model will
                 be a 2D tensor.
             - `max` means that global max pooling will be applied.
-        alpha: controls the width of the network. This is known as the
+        alpha: float, controls the width of the network. This is known as the
             depth multiplier in the MobileNetV3 paper, but the name is kept for
             consistency with MobileNetV1 in Keras.
             - If `alpha` < 1.0, proportionally decreases the number
@@ -72,8 +72,8 @@ BASE_DOCSTRING = """Instantiates the {name} architecture.
                 of filters in each layer.
             - If `alpha` = 1, default number of filters from the paper
                 are used at each layer.
-        minimalistic: in addition to large and small models this module also
-            contains so-called minimalistic models, these models have the same
+        minimalistic: in addition to large and small models, this module also
+            contains so-called minimalistic models; these models have the same
             per-layer dimensions characteristic as MobilenetV3 however, they don't
             utilize any of the advanced blocks (squeeze-and-excite units, hard-swish,
             and 5x5 convolutions). While these models are less efficient on CPU, they
@@ -81,7 +81,7 @@ BASE_DOCSTRING = """Instantiates the {name} architecture.
         dropout_rate: a float between 0 and 1 denoting the fraction of input units to
             drop, defaults to 0.2.
         classifier_activation: the activation function to use, defaults to softmax.
-        name: (Optional) name to pass to the model. Defaults to "{name}".
+        name: string, optional name to pass to the model, defaults to "{name}".
 
     Returns:
         A `keras.Model` instance.
@@ -89,16 +89,16 @@ BASE_DOCSTRING = """Instantiates the {name} architecture.
 
 
 def depth(x, divisor=8, min_value=None):
-    """Ensure that all layers have a channel number that is divisble by the `divisor`.
+    """Ensure that all layers have a channel number that is divisible by the `divisor`.
 
     Args:
-        x: input value.
-        divisor: integer, the value by which a channel number should be divisble,
+        x: integer, input value.
+        divisor: integer, the value by which a channel number should be divisible,
             defaults to 8.
         min_value: float, minimum value for the new tensor.
 
     Returns:
-        the updated value of the input.
+        the updated input scalar.
     """
 
     if min_value is None:
@@ -112,48 +112,44 @@ def depth(x, divisor=8, min_value=None):
     return new_x
 
 
-def HardSigmoid(name=None):
+def apply_hard_sigmoid(x, name=None):
     """The Hard Sigmoid function.
 
     Args:
+        x: input tensor
         name: string, layer label.
 
     Returns:
-        a function that takes an input Tensor representing a HardSigmoid layer.
+        the updated input tensor.
     """
     if name is None:
         name = f"hard_sigmoid_{backend.get_uid('hard_sigmoid')}"
 
     activation = layers.ReLU(6.0)
 
-    def apply(x):
-        return activation(x + 3.0) * (1.0 / 6.0)
+    return activation(x + 3.0) * (1.0 / 6.0)
 
-    return apply
 
-
-def HardSwish(name=None):
+def apply_hard_swish(x, name=None):
     """The Hard Swish function.
 
     Args:
+        x: input tensor
         name: string, layer label.
 
     Returns:
-        a function that takes an input Tensor representing a HardSwish layer.
+        the updated input tensor.
     """
     if name is None:
         name = f"hard_swish_{backend.get_uid('hard_swish')}"
 
-    hard_sigmoid = HardSigmoid()
     multiply_layer = layers.Multiply()
 
-    def apply(x):
-        return multiply_layer([x, hard_sigmoid(x)])
-
-    return apply
+    return multiply_layer([x, apply_hard_sigmoid(x)])
 
 
-def InvertedResBlock(
+def apply_inverted_res_block(
+    x,
     expansion,
     filters,
     kernel_size,
@@ -166,11 +162,12 @@ def InvertedResBlock(
     """An Inverted Residual Block.
 
     Args:
+        x: input tensor.
         expansion: integer, the expansion ratio, multiplied with infilters to get the
             minimum value passed to depth.
         filters: integer, number of filters for convolution layer.
-        kernel_size: integer, the kernel size for DpethWise Convolutions.
-        stride: integer, the stride length for DpethWise Convolutions.
+        kernel_size: integer, the kernel size for DepthWise Convolutions.
+        stride: integer, the stride length for DepthWise Convolutions.
         se_ratio: float, ratio for bottleneck filters. Number of bottleneck
             filters = filters * se_ratio.
         activation: the activation layer to use.
@@ -179,104 +176,86 @@ def InvertedResBlock(
         name: string, layer label.
 
     Returns:
-        a function that takes an input Tensor representing a InvertedResBlock.
+        the updated input tensor.
     """
     if name is None:
         name = f"inverted_res_block_{backend.get_uid('inverted_res_block')}"
 
-    def apply(x):
-        shortcut = x
-        prefix = "expanded_conv/"
-        infilters = backend.int_shape(x)[channel_axis]
-
-        if block_id:
-            prefix = f"expanded_conv_{block_id}"
-
-            x = layers.Conv2D(
-                depth(infilters * expansion),
-                kernel_size=1,
-                padding="same",
-                use_bias=False,
-                name=prefix + "expand",
-            )(x)
-            x = layers.BatchNormalization(
-                axis=channel_axis,
-                epsilon=1e-3,
-                momentum=0.999,
-                name=prefix + "expand/BatchNorm",
-            )(x)
-            x = activation(x)
-
-        x = layers.DepthwiseConv2D(
-            kernel_size,
-            strides=stride,
-            padding="same" if stride == 1 else "valid",
-            use_bias=False,
-            name=prefix + "depthwise",
-        )(x)
-        x = layers.BatchNormalization(
-            axis=channel_axis,
-            epsilon=1e-3,
-            momentum=0.999,
-            name=prefix + "depthwise/BatchNorm",
-        )(x)
-        x = activation(x)
+    shortcut = x
+    prefix = "expanded_conv/"
+    infilters = backend.int_shape(x)[channel_axis]
 
-        if se_ratio:
-            with custom_object_scope({"hard_sigmoid": HardSigmoid()}):
-                x = cv_layers.SqueezeAndExcite2D(
-                    filters=depth(infilters * expansion),
-                    ratio=se_ratio,
-                    squeeze_activation="relu",
-                    excite_activation="hard_sigmoid",
-                )(x)
+    if block_id:
+        prefix = f"expanded_conv_{block_id}"
 
         x = layers.Conv2D(
-            filters,
+            depth(infilters * expansion),
             kernel_size=1,
             padding="same",
             use_bias=False,
-            name=prefix + "project",
+            name=prefix + "expand",
         )(x)
         x = layers.BatchNormalization(
             axis=channel_axis,
             epsilon=1e-3,
             momentum=0.999,
-            name=prefix + "project/BatchNorm",
+            name=prefix + "expand/BatchNorm",
         )(x)
+        x = activation(x)
 
-        if stride == 1 and infilters == filters:
-            x = layers.Add(name=prefix + "Add")([shortcut, x])
+    x = layers.DepthwiseConv2D(
+        kernel_size,
+        strides=stride,
+        padding="same" if stride == 1 else "valid",
+        use_bias=False,
+        name=prefix + "depthwise",
+    )(x)
+    x = layers.BatchNormalization(
+        axis=channel_axis,
+        epsilon=1e-3,
+        momentum=0.999,
+        name=prefix + "depthwise/BatchNorm",
+    )(x)
+    x = activation(x)
 
-        return x
+    if se_ratio:
+        with custom_object_scope({"hard_sigmoid": apply_hard_sigmoid}):
+            x = cv_layers.SqueezeAndExcite2D(
+                filters=depth(infilters * expansion),
+                ratio=se_ratio,
+                squeeze_activation="relu",
+                excite_activation="hard_sigmoid",
+            )(x)
 
-    return apply
+    x = layers.Conv2D(
+        filters,
+        kernel_size=1,
+        padding="same",
+        use_bias=False,
+        name=prefix + "project",
+    )(x)
+    x = layers.BatchNormalization(
+        axis=channel_axis,
+        epsilon=1e-3,
+        momentum=0.999,
+        name=prefix + "project/BatchNorm",
+    )(x)
 
+    if stride == 1 and infilters == filters:
+        x = layers.Add(name=prefix + "Add")([shortcut, x])
 
-def MobileNetV3(
-    stack_fn,
-    last_point_ch,
-    include_rescaling,
-    include_top,
-    num_classes=None,
-    weights=None,
-    input_shape=(None, None, 3),
-    input_tensor=None,
-    pooling=None,
-    alpha=1.0,
-    minimalistic=True,
-    dropout_rate=0.2,
-    classifier_activation="softmax",
-    name="MobileNetV3",
-    **kwargs,
-):
+    return x
+
+
+@keras.utils.register_keras_serializable(package="keras_cv.models")
+class MobileNetV3(keras.Model):
     """Instantiates the MobileNetV3 architecture.
 
     References:
         - [Searching for MobileNetV3](https://arxiv.org/pdf/1905.02244.pdf) (ICCV 2019)
         - [Based on the Original keras.applications MobileNetv3](https://github.com/keras-team/keras/blob/master/keras/applications/mobilenet_v3.py)
 
-    This function returns a Keras MobileNetV3 model.
+    This class represents a Keras MobileNetV3 model.
 
     For transfer learning use cases, make sure to read the [guide to transfer
         learning & fine-tuning](https://keras.io/guides/transfer_learning/).
@@ -284,19 +263,19 @@ def MobileNetV3(
     Args:
         stack_fn: a function that returns tensors passed through Inverted
             Residual Blocks.
-        last_point_ch: the number of filters for the convolution layer.
-        include_rescaling: whether or not to Rescale the inputs.If set to True,
+        last_point_ch: integer, the number of filters for the convolution layer.
+        include_rescaling: bool, whether to rescale the inputs. If set to True,
             inputs will be passed through a `Rescaling(scale=1 / 255)`
-            layer, defaults to True.
-        include_top: whether to include the fully-connected layer at the top of the
-            network.  If provided, `num_classes` must be provided.
-        num_classes: optional number of classes to classify images into, only to be
+            layer.
+        include_top: bool, whether to include the fully-connected layer at the top of the
+            network. If provided, `num_classes` must be provided.
+        num_classes: optional number of classes to classify images into. Only to be
             specified if `include_top` is True, and if no `weights` argument is
             specified.
-        weights: one of `None` (random initialization), or a pretrained weight file
+        weights: one of `None` (random initialization) or a pre-trained weight file
             path.
         input_shape: optional shape tuple, defaults to (None, None, 3).
-        input_tensor: optional Keras tensor (i.e. output of `layers.Input()`)
+        input_tensor: optional Keras tensor (i.e., output of `layers.Input()`)
             to use as image input for the model.
         pooling: optional pooling mode for feature extraction
             when `include_top` is `False`.
@@ -306,7 +285,7 @@ def MobileNetV3(
                 of the last convolutional block, and thus the output of the model will
                 be a 2D tensor.
             - `max` means that global max pooling will be applied.
-        alpha: controls the width of the network. This is known as the
+        alpha: float, controls the width of the network. This is known as the
             depth multiplier in the MobileNetV3 paper, but the name is kept for
             consistency with MobileNetV1 in Keras.
             - If `alpha` < 1.0, proportionally decreases the number
@@ -315,8 +294,8 @@ def MobileNetV3(
                 of filters in each layer.
             - If `alpha` = 1, default number of filters from the paper
                 are used at each layer.
-        minimalistic: in addition to large and small models this module also
-            contains so-called minimalistic models, these models have the same
+        minimalistic: in addition to large and small models, this module also
+            contains so-called minimalistic models; these models have the same
             per-layer dimensions characteristic as MobilenetV3 however, they don't
             utilize any of the advanced blocks (squeeze-and-excite units, hard-swish,
             and 5x5 convolutions). While these models are less efficient on CPU, they
@@ -324,109 +303,167 @@ def MobileNetV3(
         dropout_rate: a float between 0 and 1 denoting the fraction of input units to
             drop, defaults to 0.2.
         classifier_activation: the activation function to use, defaults to softmax.
-
-        name: (Optional) name to pass to the model. Defaults to "MobileNetV3".
+        name: string, optional name to pass to the model, defaults to "MobileNetV3".
+        **kwargs: Pass-through keyword arguments to `tf.keras.Model`.
 
     Returns:
         A `keras.Model` instance.
 
     Raises:
-        ValueError: if `weights` represents an invalid path to weights file and is not
+        ValueError: if `weights` represents an invalid path to the weights file and is not
             None.
         ValueError: if `include_top` is True and `num_classes` is not specified.
     """
-    if weights and not tf.io.gfile.exists(weights):
-        raise ValueError(
-            "The `weights` argument should be either "
-            "`None` or the path to the weights file to be loaded. "
-            f"Weights file not found at location: {weights}"
-        )
-
-    if include_top and not num_classes:
-        raise ValueError(
-            "If `include_top` is True, "
-            "you should specify `num_classes`. "
-            f"Received: num_classes={num_classes}"
-        )
-
-    if minimalistic:
-        kernel = 3
-        activation = layers.ReLU()
-        se_ratio = None
-    else:
-        kernel = 5
-        activation = HardSwish()
-        se_ratio = 0.25
-
-    inputs = utils.parse_model_inputs(input_shape, input_tensor)
-
-    x = inputs
 
-    if include_rescaling:
-        x = layers.Rescaling(scale=1 / 255)(x)
-
-    x = layers.Conv2D(
-        16,
-        kernel_size=3,
-        strides=(2, 2),
-        padding="same",
-        use_bias=False,
-        name="Conv",
-    )(x)
-    x = layers.BatchNormalization(
-        axis=channel_axis, epsilon=1e-3, momentum=0.999, name="Conv/BatchNorm"
-    )(x)
-    x = activation(x)
-
-    x = stack_fn(x, kernel, activation, se_ratio)
-
-    last_conv_ch = depth(backend.int_shape(x)[channel_axis] * 6)
+    def __init__(
+        self,
+        stack_fn,
+        last_point_ch,
+        include_rescaling,
+        include_top,
+        num_classes=None,
+        weights=None,
+        input_shape=(None, None, 3),
+        input_tensor=None,
+        pooling=None,
+        alpha=1.0,
+        minimalistic=True,
+        dropout_rate=0.2,
+        classifier_activation="softmax",
+        name="MobileNetV3",
+        **kwargs,
+    ):
+        if weights and not tf.io.gfile.exists(weights):
+            raise ValueError(
+                "The `weights` argument should be either "
+                "`None` or the path to the weights file to be loaded. "
+                f"Weights file not found at location: {weights}"
+            )
+
+        if include_top and not num_classes:
+            raise ValueError(
+                "If `include_top` is True, "
+                "you should specify `num_classes`. "
+                f"Received: num_classes={num_classes}"
+            )
+
+        if minimalistic:
+            kernel = 3
+            activation = layers.ReLU()
+            se_ratio = None
+        else:
+            kernel = 5
+            activation = apply_hard_swish
+            se_ratio = 0.25
+
+        inputs = utils.parse_model_inputs(input_shape, input_tensor)
+
+        x = inputs
+
+        if include_rescaling:
+            x = layers.Rescaling(scale=1 / 255)(x)
 
-    # if the width multiplier is greater than 1 we
-    # increase the number of output channels
-    if alpha > 1.0:
-        last_point_ch = depth(last_point_ch * alpha)
-    x = layers.Conv2D(
-        last_conv_ch,
-        kernel_size=1,
-        padding="same",
-        use_bias=False,
-        name="Conv_1",
-    )(x)
-    x = layers.BatchNormalization(
-        axis=channel_axis, epsilon=1e-3, momentum=0.999, name="Conv_1/BatchNorm"
-    )(x)
-    x = activation(x)
-    if include_top:
-        x = layers.GlobalAveragePooling2D(keepdims=True)(x)
         x = layers.Conv2D(
-            last_point_ch,
-            kernel_size=1,
+            16,
+            kernel_size=3,
+            strides=(2, 2),
             padding="same",
-            use_bias=True,
-            name="Conv_2",
+            use_bias=False,
+            name="Conv",
+        )(x)
+        x = layers.BatchNormalization(
+            axis=channel_axis,
+            epsilon=1e-3,
+            momentum=0.999,
+            name="Conv/BatchNorm",
         )(x)
         x = activation(x)
 
-        if dropout_rate > 0:
-            x = layers.Dropout(dropout_rate)(x)
+        x = stack_fn(x, kernel, activation, se_ratio)
+
+        last_conv_ch = depth(backend.int_shape(x)[channel_axis] * 6)
+
+        # if the width multiplier is greater than 1 we
+        # increase the number of output channels
+        if alpha > 1.0:
+            last_point_ch = depth(last_point_ch * alpha)
         x = layers.Conv2D(
-            num_classes, kernel_size=1, padding="same", name="Logits"
+            last_conv_ch,
+            kernel_size=1,
+            padding="same",
+            use_bias=False,
+            name="Conv_1",
         )(x)
-        x = layers.Flatten()(x)
-        x = layers.Activation(
-            activation=classifier_activation, name="Predictions"
+        x = layers.BatchNormalization(
+            axis=channel_axis,
+            epsilon=1e-3,
+            momentum=0.999,
+            name="Conv_1/BatchNorm",
         )(x)
-    elif pooling == "avg":
-        x = layers.GlobalAveragePooling2D(name="avg_pool")(x)
-    elif pooling == "max":
-        x = layers.GlobalMaxPooling2D(name="max_pool")(x)
-
-    model = keras.Model(inputs, x, name=name, **kwargs)
+        x = activation(x)
+        if include_top:
+            x = layers.GlobalAveragePooling2D(keepdims=True)(x)
+            x = layers.Conv2D(
+                last_point_ch,
+                kernel_size=1,
+                padding="same",
+                use_bias=True,
+                name="Conv_2",
+            )(x)
+            x = activation(x)
 
-    if weights is not None:
-        model.load_weights(weights)
-    return model
+            if dropout_rate > 0:
+                x = layers.Dropout(dropout_rate)(x)
+            x = layers.Conv2D(
+                num_classes, kernel_size=1, padding="same", name="Logits"
+            )(x)
+            x = layers.Flatten()(x)
+            x = layers.Activation(
+                activation=classifier_activation, name="Predictions"
+            )(x)
+        elif pooling == "avg":
+            x = layers.GlobalAveragePooling2D(name="avg_pool")(x)
+        elif pooling == "max":
+            x = layers.GlobalMaxPooling2D(name="max_pool")(x)
+
+        super().__init__(inputs=inputs, outputs=x, name=name, **kwargs)
+
+        if weights is not None:
+            self.load_weights(weights)
+
+        self.stack_fn = stack_fn
+        self.last_point_ch = last_point_ch
+        self.include_rescaling = include_rescaling
+        self.include_top = include_top
+        self.num_classes = num_classes
+        self.input_tensor = input_tensor
+        self.pooling = pooling
+        self.alpha = alpha
+        self.minimalistic = minimalistic
+        self.dropout_rate = dropout_rate
+        self.classifier_activation = classifier_activation
+
+    def get_config(self):
+        return {
+            "stack_fn": self.stack_fn,
+            "last_point_ch": self.last_point_ch,
+            "include_rescaling": self.include_rescaling,
+            "include_top": self.include_top,
+            "num_classes": self.num_classes,
+            "input_shape": self.input_shape[1:],
+            "input_tensor": self.input_tensor,
+            "pooling": self.pooling,
+            "alpha": self.alpha,
+            "minimalistic": self.minimalistic,
+            "dropout_rate": self.dropout_rate,
+            "classifier_activation": self.classifier_activation,
+            "name": self.name,
+            "trainable": self.trainable,
+        }
+
+    @classmethod
+    def from_config(cls, config):
+        return cls(**config)
 
 
 def MobileNetV3Small(
@@ -446,39 +483,39 @@ def MobileNetV3Small(
     **kwargs,
 ):
     def stack_fn(x, kernel, activation, se_ratio):
-        x = InvertedResBlock(
-            1, depth(16 * alpha), 3, 2, se_ratio, layers.ReLU(), 0
-        )(x)
-        x = InvertedResBlock(
-            72.0 / 16, depth(24 * alpha), 3, 2, None, layers.ReLU(), 1
-        )(x)
-        x = InvertedResBlock(
-            88.0 / 24, depth(24 * alpha), 3, 1, None, layers.ReLU(), 2
-        )(x)
-        x = InvertedResBlock(
-            4, depth(40 * alpha), kernel, 2, se_ratio, activation, 3
-        )(x)
-        x = InvertedResBlock(
-            6, depth(40 * alpha), kernel, 1, se_ratio, activation, 4
-        )(x)
-        x = InvertedResBlock(
-            6, depth(40 * alpha), kernel, 1, se_ratio, activation, 5
-        )(x)
-        x = InvertedResBlock(
-            3, depth(48 * alpha), kernel, 1, se_ratio, activation, 6
-        )(x)
-        x = InvertedResBlock(
-            3, depth(48 * alpha), kernel, 1, se_ratio, activation, 7
-        )(x)
-        x = InvertedResBlock(
-            6, depth(96 * alpha), kernel, 2, se_ratio, activation, 8
-        )(x)
-        x = InvertedResBlock(
-            6, depth(96 * alpha), kernel, 1, se_ratio, activation, 9
-        )(x)
-        x = InvertedResBlock(
-            6, depth(96 * alpha), kernel, 1, se_ratio, activation, 10
-        )(x)
+        x = apply_inverted_res_block(
+            x, 1, depth(16 * alpha), 3, 2, se_ratio, layers.ReLU(), 0
+        )
+        x = apply_inverted_res_block(
+            x, 72.0 / 16, depth(24 * alpha), 3, 2, None, layers.ReLU(), 1
+        )
+        x = apply_inverted_res_block(
+            x, 88.0 / 24, depth(24 * alpha), 3, 1, None, layers.ReLU(), 2
+        )
+        x = apply_inverted_res_block(
+            x, 4, depth(40 * alpha), kernel, 2, se_ratio, activation, 3
+        )
+        x = apply_inverted_res_block(
+            x, 6, depth(40 * alpha), kernel, 1, se_ratio, activation, 4
+        )
+        x = apply_inverted_res_block(
+            x, 6, depth(40 * alpha), kernel, 1, se_ratio, activation, 5
+        )
+        x = apply_inverted_res_block(
+            x, 3, depth(48 * alpha), kernel, 1, se_ratio, activation, 6
+        )
+        x = apply_inverted_res_block(
+            x, 3, depth(48 * alpha), kernel, 1, se_ratio, activation, 7
+        )
+        x = apply_inverted_res_block(
+            x, 6, depth(96 * alpha), kernel, 2, se_ratio, activation, 8
+        )
+        x = apply_inverted_res_block(
+            x, 6, depth(96 * alpha), kernel, 1, se_ratio, activation, 9
+        )
+        x = apply_inverted_res_block(
+            x, 6, depth(96 * alpha), kernel, 1, se_ratio, activation, 10
+        )
         return x
 
     return MobileNetV3(
@@ -517,49 +554,51 @@ def MobileNetV3Large(
     **kwargs,
 ):
     def stack_fn(x, kernel, activation, se_ratio):
-        x = InvertedResBlock(
-            1, depth(16 * alpha), 3, 1, None, layers.ReLU(), 0
-        )(x)
-        x = InvertedResBlock(
-            4, depth(24 * alpha), 3, 2, None, layers.ReLU(), 1
-        )(x)
-        x = InvertedResBlock(
-            3, depth(24 * alpha), 3, 1, None, layers.ReLU(), 2
-        )(x)
-        x = InvertedResBlock(
-            3, depth(40 * alpha), kernel, 2, se_ratio, layers.ReLU(), 3
-        )(x)
-        x = InvertedResBlock(
-            3, depth(40 * alpha), kernel, 1, se_ratio, layers.ReLU(), 4
-        )(x)
-        x = InvertedResBlock(
-            3, depth(40 * alpha), kernel, 1, se_ratio, layers.ReLU(), 5
-        )(x)
-        x = InvertedResBlock(6, depth(80 * alpha), 3, 2, None, activation, 6)(x)
-        x = InvertedResBlock(2.5, depth(80 * alpha), 3, 1, None, activation, 7)(
-            x
+        x = apply_inverted_res_block(
+            x, 1, depth(16 * alpha), 3, 1, None, layers.ReLU(), 0
         )
-        x = InvertedResBlock(2.3, depth(80 * alpha), 3, 1, None, activation, 8)(
-            x
+        x = apply_inverted_res_block(
+            x, 4, depth(24 * alpha), 3, 2, None, layers.ReLU(), 1
         )
-        x = InvertedResBlock(2.3, depth(80 * alpha), 3, 1, None, activation, 9)(
-            x
+        x = apply_inverted_res_block(
+            x, 3, depth(24 * alpha), 3, 1, None, layers.ReLU(), 2
+        )
+        x = apply_inverted_res_block(
+            x, 3, depth(40 * alpha), kernel, 2, se_ratio, layers.ReLU(), 3
+        )
+        x = apply_inverted_res_block(
+            x, 3, depth(40 * alpha), kernel, 1, se_ratio, layers.ReLU(), 4
+        )
+        x = apply_inverted_res_block(
+            x, 3, depth(40 * alpha), kernel, 1, se_ratio, layers.ReLU(), 5
+        )
+        x = apply_inverted_res_block(
+            x, 6, depth(80 * alpha), 3, 2, None, activation, 6
+        )
+        x = apply_inverted_res_block(
+            x, 2.5, depth(80 * alpha), 3, 1, None, activation, 7
+        )
+        x = apply_inverted_res_block(
+            x, 2.3, depth(80 * alpha), 3, 1, None, activation, 8
+        )
+        x = apply_inverted_res_block(
+            x, 2.3, depth(80 * alpha), 3, 1, None, activation, 9
+        )
+        x = apply_inverted_res_block(
+            x, 6, depth(112 * alpha), 3, 1, se_ratio, activation, 10
+        )
+        x = apply_inverted_res_block(
+            x, 6, depth(112 * alpha), 3, 1, se_ratio, activation, 11
+        )
+        x = apply_inverted_res_block(
+            x, 6, depth(160 * alpha), kernel, 2, se_ratio, activation, 12
+        )
+        x = apply_inverted_res_block(
+            x, 6, depth(160 * alpha), kernel, 1, se_ratio, activation, 13
+        )
+        x = apply_inverted_res_block(
+            x, 6, depth(160 * alpha), kernel, 1, se_ratio, activation, 14
         )
-        x = InvertedResBlock(
-            6, depth(112 * alpha), 3, 1, se_ratio, activation, 10
-        )(x)
-        x = InvertedResBlock(
-            6, depth(112 * alpha), 3, 1, se_ratio, activation, 11
-        )(x)
-        x = InvertedResBlock(
-            6, depth(160 * alpha), kernel, 2, se_ratio, activation, 12
-        )(x)
-        x = InvertedResBlock(
-            6, depth(160 * alpha), kernel, 1, se_ratio, activation, 13
-        )(x)
-        x = InvertedResBlock(
-            6, depth(160 * alpha), kernel, 1, se_ratio, activation, 14
-        )(x)
         return x
 
     return MobileNetV3(
diff --git a/keras_cv/models/object_detection_3d/center_pillar.py b/keras_cv/models/object_detection_3d/center_pillar.py
index 6aed94b..54dcb4c 100644
--- a/keras_cv/models/object_detection_3d/center_pillar.py
+++ b/keras_cv/models/object_detection_3d/center_pillar.py
@@ -158,7 +158,12 @@ class MultiHeadCenterPillar(tf.keras.Model):
         self._prediction_decoder = prediction_decoder
         self._head_names = self._multiclass_head._head_names
 
-    def call(self, point_xyz, point_feature, point_mask, training=None):
+    def call(self, input_dict, training=None):
+        point_xyz, point_feature, point_mask = (
+            input_dict["point_xyz"],
+            input_dict["point_feature"],
+            input_dict["point_mask"],
+        )
         voxel_feature = self._voxelization_layer(
             point_xyz, point_feature, point_mask, training=training
         )
@@ -169,9 +174,11 @@ class MultiHeadCenterPillar(tf.keras.Model):
         # returns dict {"class_1": concat_pred_1, "class_2": concat_pred_2}
         return predictions
 
-    def compute_loss(
-        self, predictions, box_dict, heatmap_dict, top_k_index_dict
-    ):
+    def compute_loss(self, predictions, targets):
+        box_dict = targets["box_3d"]
+        heatmap_dict = targets["heatmap"]
+        top_k_index_dict = targets["top_k_index"]
+
         y_pred = {}
         y_true = {}
         sample_weight = {}
@@ -200,25 +207,8 @@ class MultiHeadCenterPillar(tf.keras.Model):
 
     def train_step(self, data):
         x, y, sample_weight = tf.keras.utils.unpack_x_y_sample_weight(data)
-        box_3d_dict = y["box_3d"]
-        heatmap_dict = y["heatmap"]
-        top_k_index_dict = y["top_k_index"]
-        losses = []
         with tf.GradientTape() as tape:
-            predictions = self(
-                x["point_xyz"],
-                x["point_feature"],
-                x["point_mask"],
-                training=True,
-            )
-            losses.append(
-                self.compute_loss(
-                    predictions, box_3d_dict, heatmap_dict, top_k_index_dict
-                )
-            )
-            if self.weight_decay:
-                for var in self.trainable_variables:
-                    losses.append(self.weight_decay * tf.nn.l2_loss(var))
-            total_loss = tf.math.add_n(losses)
-        self.optimizer.minimize(total_loss, self.trainable_variables, tape=tape)
+            predictions = self(x, training=True)
+            loss = self.compute_loss(predictions, y)
+        self.optimizer.minimize(loss, self.trainable_variables, tape=tape)
         return self.compute_metrics({}, {}, {}, sample_weight={})


Test patch:
diff --git a/keras_cv/models/object_detection_3d/center_pillar_test.py b/keras_cv/models/object_detection_3d/center_pillar_test.py
index 3d64fc7..f9b77e3 100644
--- a/keras_cv/models/object_detection_3d/center_pillar_test.py
+++ b/keras_cv/models/object_detection_3d/center_pillar_test.py
@@ -93,7 +93,14 @@ class CenterPillarTest(tf.test.TestCase):
         point_xyz = tf.random.normal([2, 1000, 3])
         point_feature = tf.random.normal([2, 1000, 4])
         point_mask = tf.constant(True, shape=[2, 1000])
-        outputs = model(point_xyz, point_feature, point_mask, training=True)
+        outputs = model(
+            {
+                "point_xyz": point_xyz,
+                "point_feature": point_feature,
+                "point_mask": point_mask,
+            },
+            training=True,
+        )
         self.assertEqual(outputs["class_1"].shape, [2, 400, 400, 12])
         self.assertEqual(outputs["class_2"].shape, [2, 400, 400, 12])
 
@@ -129,7 +136,14 @@ class CenterPillarTest(tf.test.TestCase):
         point_xyz = tf.random.normal([2, 1000, 3])
         point_feature = tf.random.normal([2, 1000, 4])
         point_mask = tf.constant(True, shape=[2, 1000])
-        outputs = model(point_xyz, point_feature, point_mask, training=False)
+        outputs = model(
+            {
+                "point_xyz": point_xyz,
+                "point_feature": point_feature,
+                "point_mask": point_mask,
+            },
+            training=False,
+        )
         # max number boxes is 3
         self.assertEqual(outputs["class_1"][0].shape, [2, 3, 7])
         self.assertEqual(outputs["class_1"][1].shape, [2, 3])