Repository: keras-team/keras-cv

Issue:
Vectorize ChannelShuffle
The PR contributing this should be based on:

https://github.com/keras-team/keras-cv/pull/1394

Discussion:
james77777778: In this PR, I find that channel shuffle cannot process independently.
Need some time to survey but it could be difficult because `tf.random.shuffle` cannot apply on batched inputs.
LukeWood: > In this PR, I find that channel shuffle cannot process independently. Need some time to survey but it could be difficult because `tf.random.shuffle` cannot apply on batched inputs.

can we roll our own version instead of using `tf.random.shuffle`?  I believe this could be done with `tf.gather_nd()` and `tf.random.uniform()`.
james77777778: > > In this PR, I find that channel shuffle cannot process independently. Need some time to survey but it could be difficult because `tf.random.shuffle` cannot apply on batched inputs.
> 
> can we roll our own version instead of using `tf.random.shuffle`? I believe this could be done with `tf.gather_nd()` and `tf.random.uniform()`.

Thanks for the keywords.
I have tested tf.TensorArray (as a output list) with similar implementation from previous one but the performance decreased a lot...

I'm going to try `tf.gather_nd()` and `tf.random.uniform()`
bhack: I think we can do:

```python
num_channels = x.shape.as_list()[-1]

    # Generate a random permutation of the channel indices
    shuffled_indices = tf.random.shuffle(tf.range(num_channels), seed=seed)

    # Permute the channels of the input tensor according to the random indices
    shuffled_x = tf.gather(x, shuffled_indices, axis=-1)
```
james77777778: @LukeWood @bhack 
Thanks for your guidance

Much of the difficulty comes from how to independently shuffle indices in batch.

I adopt the `rand+argsort` trick based on
https://stackoverflow.com/a/55317373

```python
    def get_random_transformation_batch(self, batch_size, **kwargs):
        # get batched shuffled indices
        # for example: batch_size=2; self.group=5
        # indices = [
        #     [0, 2, 3, 4, 1],
        #     [4, 1, 0, 2, 3]
        # ]
        indices_distribution = tf.random.uniform(
            (batch_size, self.groups), seed=self.seed
        )
        indices = tf.argsort(indices_distribution, axis=-1)
        return indices
```
and then gather by
```python
    def augment_images(self, images, transformations, **kwargs):
        ......
        indices = transformations

        # append batch indexes next to shuffled indices
        batch_indexs = tf.reshape(
            tf.repeat(tf.range(batch_size), self.groups),
            (batch_size, self.groups),
        )
        indices = tf.stack([batch_indexs, indices], axis=-1)

        channels_per_group = num_channels // self.groups
        images = tf.reshape(
            images, [batch_size, height, width, self.groups, channels_per_group]
        )
        images = tf.transpose(images, perm=[0, 3, 1, 2, 4])
        images = tf.gather_nd(images, indices=indices)
        images = tf.transpose(images, perm=[0, 2, 3, 4, 1])
        images = tf.reshape(images, [batch_size, height, width, num_channels])
        return images
```

All unit tests passed but it is difficult to check the consistency with non-vectorized implementation because the randomness comes from different sources (`tf.random.shuffle`, `tf.random.uniform`).
I just removed the test in benchmark script.

Also a independence test is added.

The performance result:
![comparison](https://user-images.githubusercontent.com/20734616/220811027-cfd0f737-2b37-49f7-9a14-cb4092fd461e.png)

bhack: > Much of the difficulty comes from how to independently shuffle indices in batch.

If I remember correctly is `tf.random.experimental.stateless_shuffle`
james77777778: > > Much of the difficulty comes from how to independently shuffle indices in batch.
> 
> If I remember correctly is `tf.random.experimental.stateless_shuffle`

I think `tf.random.experimental.stateless_shuffle` acts the same as `tf.random.shuffle`
So they can only shuffle tensor along 1 axis?

In ChannelShuffle, we want to shuffle channels independently in each image like:
idx=1 => shuffled_idx = [1, 2, 0]
idx=2 => shuffled_idx = [0, 2, 1]
...

It is achievable to loop over each image with the operations you mentioned.
Pseudo code like:
```python
outputs = tf.TensorArray(...)
for i in range(batch_size):
    num_channels = images[i].shape.as_list()[-1]
    shuffled_indices = tf.random.shuffle(tf.range(num_channels), seed=seed)
    shuffled_x = tf.gather(images[i], shuffled_indices, axis=-1)
    outputs.write(i, shuffled_x)
outputs.stack()
```
But it is too slow that it is even slower than non-vectorized version in eager mode.

I cannot figure out how to vectorize this with `tf.random.shuffle` or `tf.random.experimental.stateless_shuffle`. Please enlighten me if I missed something.
bhack: > In ChannelShuffle, we want to shuffle channels independently in each image like:
idx=1 => shuffled_idx = [1, 2, 0]
idx=2 => shuffled_idx = [0, 2, 1]

Something like?
https://colab.research.google.com/gist/bhack/65760a7896ee9f6ec450ebbe8bd32ca2/exp_random_shuffle.ipynb
james77777778: > > In ChannelShuffle, we want to shuffle channels independently in each image like:
> > idx=1 => shuffled_idx = [1, 2, 0]
> > idx=2 => shuffled_idx = [0, 2, 1]
> 
> Something like this? https://colab.research.google.com/gist/bhack/65760a7896ee9f6ec450ebbe8bd32ca2/exp_random_shuffle.ipynb

https://colab.research.google.com/drive/1b-tZErA-jCv6SNurdprdOYc6dLFr0ZSH?usp=sharing
I manually set the values to show that
- only 1 shuffled indices for all batch
- outputs are shuffled by same indices and not independently across batch
bhack: Ok, so you meant within the batch. Sorry I was confused on what kind of randomization policy we finally selected for the new vectorizing refactoring (mandatory within the batch randomization or not)


james77777778: Kindly ping @LukeWood @ianstenbit 

The idea of vectorizing please see https://github.com/keras-team/keras-cv/pull/1433#issuecomment-1441154826

The current implementation shuffles channel (by groups) independently for each image in the batch
ianstenbit: @LukeWood PTAL -- what's the decision re: independent augmentation in batches?
LukeWood: All augmentation should be done independently on a per-image basis.
LukeWood: /gcbrun

Thanks you @james77777778 !!!
LukeWood: Sorry one more thing:  Can add a test like the one in https://github.com/keras-team/keras-cv/pull/1480/files#diff-1c771135874332a25b437da4a5706faea85ff719331dfe7beb255403a1a66477 
james77777778: > Sorry one more thing: Can add a test like the one in https://github.com/keras-team/keras-cv/pull/1480/files#diff-1c771135874332a25b437da4a5706faea85ff719331dfe7beb255403a1a66477

Hi @LukeWood 
It should be already added at 
https://github.com/keras-team/keras-cv/pull/1433/files#diff-67ea20e78d28e97daedb7b45698d50914d1ffef26cc62311c855606d6d171fc0R98-R105

Patch:
diff --git a/.devcontainer/Dockerfile b/.devcontainer/Dockerfile
index 0229525..0e11469 100644
--- a/.devcontainer/Dockerfile
+++ b/.devcontainer/Dockerfile
@@ -1,16 +1,2 @@
-ARG VERSION
-
-FROM tensorflow/tensorflow:${VERSION}
-
-RUN apt-get update && apt-get install -y \
-  locales \
-  && rm -rf /var/lib/apt/lists/*
-
-RUN pip install \
-  black[jupyter] \
-  flake8 \
-  isort \
-  pytest \
-  regex \
-  tensorflow_datasets \
-  pycocotools
+FROM mcr.microsoft.com/vscode/devcontainers/python:3.8
+COPY setup.sh /setup.sh
diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json
index 38f5507..3f5b921 100644
--- a/.devcontainer/devcontainer.json
+++ b/.devcontainer/devcontainer.json
@@ -2,45 +2,48 @@
 	"name": "Keras-cv",
 	"build": {
 		"dockerfile": "Dockerfile",
-		"args": { 
+		"args": {
 			"VERSION": "2.11.0"
 			// Uncomment this if GPU support is required
 			// "VERSION": "2.11.0-gpu",
 		}
 	},
-
-	"settings": {
-		"python.pythonPath": "/usr/bin/python3",
-		"python.linting.enabled": true,
-		"python.linting.flake8Enabled": true,
-		"python.testing.pytestEnabled":true,
-		"python.editor.defaultFormatter": "ms-python.black-formatter",
-		"python.editor.formatOnSave": true,
-		"python.editor.codeActionsOnSave": {
-			"source.organizeImports": true
+	"customizations": {
+		"vscode": {
+			"settings": {
+				"python.linting.enabled": true,
+				"python.linting.flake8Enabled": true,
+				"python.linting.pylintEnabled": false,
+				"python.testing.pytestEnabled": true,
+				"editor.formatOnSave": true,
+				"editor.codeActionsOnSave": {
+					"source.organizeImports": true
+				},
+				"[python]": {
+					"editor.defaultFormatter": "ms-python.black-formatter"
+				},
+				"editor.rulers": [
+					80
+				]
+			},
+			"extensions": [
+				"ms-python.python",
+				"ms-python.isort",
+				"ms-python.flake8",
+				"ms-python.black-formatter"
+			]
 		}
 	},
-
-	"extensions": [
-		"ms-python.python",
-		"ms-python.isort",
-		"ms-python.black-formatter"
-	],
-
 	"features": {
-		"git": {
-			"version": "os-provided"
-		}
+		"ghcr.io/devcontainers/features/github-cli:1": {}
 	},
-
 	// TODO: Improve to allow dynamic runArgs, see microsoft/vscode-remote-release#3972
 	// Uncomment this if GPU support is required
 	// "runArgs": [
-    	//     "--gpus=all"
+	//     "--gpus=all"
 	// ],
-
-	"onCreateCommand": "locale-gen \"en_US.UTF-8\""
+	"onCreateCommand": "locale-gen \"en_US.UTF-8\"",
 	// Optional: install pre-commit hooks
 	// "postCreateCommand": "git config core.hooksPath .github/.githooks"
-
-}
+	"postCreateCommand": "sh /setup.sh"
+}
\ No newline at end of file
diff --git a/.devcontainer/setup.sh b/.devcontainer/setup.sh
new file mode 100644
index 0000000..6f2daf0
--- /dev/null
+++ b/.devcontainer/setup.sh
@@ -0,0 +1,2 @@
+sudo pip install --upgrade pip
+sudo pip install -e ".[tests]"
diff --git a/README.md b/README.md
index 841b450..463d22c 100644
--- a/README.md
+++ b/README.md
@@ -148,7 +148,7 @@ a simple `1/255` rescaling layer.
 This can be seen in all KerasCV training pipelines and code examples.
 
 ## Custom Ops
-Note that in some the 3D Object Detection layers, custom TF ops are used. The
+Note that in some of the 3D Object Detection layers, custom TF ops are used. The
 binaries for these ops are not shipped in our PyPi package in order to keep our
 wheels pure-Python.
 
diff --git a/benchmarks/vectorized_channel_shuffle.py b/benchmarks/vectorized_channel_shuffle.py
new file mode 100644
index 0000000..1cc756e
--- /dev/null
+++ b/benchmarks/vectorized_channel_shuffle.py
@@ -0,0 +1,215 @@
+# Copyright 2023 The KerasCV Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import time
+
+import matplotlib.pyplot as plt
+import numpy as np
+import tensorflow as tf
+
+from keras_cv.layers import ChannelShuffle
+from keras_cv.layers.preprocessing.base_image_augmentation_layer import (
+    BaseImageAugmentationLayer,
+)
+
+
+class OldChannelShuffle(BaseImageAugmentationLayer):
+    """Shuffle channels of an input image.
+
+    Input shape:
+        The expected images should be [0-255] pixel ranges.
+        3D (unbatched) or 4D (batched) tensor with shape:
+        `(..., height, width, channels)`, in `"channels_last"` format
+
+    Output shape:
+        3D (unbatched) or 4D (batched) tensor with shape:
+        `(..., height, width, channels)`, in `"channels_last"` format
+
+    Args:
+        groups: Number of groups to divide the input channels. Default 3.
+        seed: Integer. Used to create a random seed.
+
+    Call arguments:
+        inputs: Tensor representing images of shape
+            `(batch_size, width, height, channels)`, with dtype tf.float32 / tf.uint8,
+            ` or (width, height, channels)`, with dtype tf.float32 / tf.uint8
+        training: A boolean argument that determines whether the call should be run
+            in inference mode or training mode. Default: True.
+
+    Usage:
+    ```python
+    (images, labels), _ = tf.keras.datasets.cifar10.load_data()
+    channel_shuffle = keras_cv.layers.ChannelShuffle()
+    augmented_images = channel_shuffle(images)
+    ```
+    """
+
+    def __init__(self, groups=3, seed=None, **kwargs):
+        super().__init__(seed=seed, **kwargs)
+        self.groups = groups
+        self.seed = seed
+
+    def augment_image(self, image, transformation=None, **kwargs):
+        shape = tf.shape(image)
+        height, width = shape[0], shape[1]
+        num_channels = image.shape[2]
+
+        if not num_channels % self.groups == 0:
+            raise ValueError(
+                "The number of input channels should be "
+                "divisible by the number of groups."
+                f"Received: channels={num_channels}, groups={self.groups}"
+            )
+
+        channels_per_group = num_channels // self.groups
+        image = tf.reshape(
+            image, [height, width, self.groups, channels_per_group]
+        )
+        image = tf.transpose(image, perm=[2, 0, 1, 3])
+        image = tf.random.shuffle(image, seed=self.seed)
+        image = tf.transpose(image, perm=[1, 2, 3, 0])
+        image = tf.reshape(image, [height, width, num_channels])
+
+        return image
+
+    def augment_bounding_boxes(self, bounding_boxes, **kwargs):
+        return bounding_boxes
+
+    def augment_label(self, label, transformation=None, **kwargs):
+        return label
+
+    def augment_segmentation_mask(
+        self, segmentation_mask, transformation, **kwargs
+    ):
+        return segmentation_mask
+
+    def get_config(self):
+        config = super().get_config()
+        config.update({"groups": self.groups, "seed": self.seed})
+        return config
+
+    def compute_output_shape(self, input_shape):
+        return input_shape
+
+
+class ChannelShuffleTest(tf.test.TestCase):
+    def test_consistency_with_old_impl(self):
+        image_shape = (1, 32, 32, 3)
+        groups = 3
+        fixed_seed = 2023  # magic number
+        image = tf.random.uniform(shape=image_shape)
+
+        layer = ChannelShuffle(groups=groups, seed=fixed_seed)
+        old_layer = OldChannelShuffle(groups=groups, seed=fixed_seed)
+
+        output = layer(image)
+        old_output = old_layer(image)
+
+        self.assertNotAllClose(image, output)
+        self.assertAllClose(old_output, output)
+
+
+if __name__ == "__main__":
+    # Run benchmark
+    (x_train, _), _ = tf.keras.datasets.cifar10.load_data()
+    x_train = x_train.astype(np.float32)
+
+    num_images = [1000, 2000, 3000, 4000, 5000, 10000]
+    results = {}
+    aug_candidates = [ChannelShuffle, OldChannelShuffle]
+    aug_args = {"groups": 3}
+
+    for aug in aug_candidates:
+        # Eager Mode
+        c = aug.__name__
+        layer = aug(**aug_args)
+        runtimes = []
+        print(f"Timing {c}")
+
+        for n_images in num_images:
+            # warmup
+            layer(x_train[:n_images])
+
+            t0 = time.time()
+            r1 = layer(x_train[:n_images])
+            t1 = time.time()
+            runtimes.append(t1 - t0)
+            print(f"Runtime for {c}, n_images={n_images}: {t1-t0}")
+        results[c] = runtimes
+
+        # Graph Mode
+        c = aug.__name__ + " Graph Mode"
+        layer = aug(**aug_args)
+
+        @tf.function()
+        def apply_aug(inputs):
+            return layer(inputs)
+
+        runtimes = []
+        print(f"Timing {c}")
+
+        for n_images in num_images:
+            # warmup
+            apply_aug(x_train[:n_images])
+
+            t0 = time.time()
+            r1 = apply_aug(x_train[:n_images])
+            t1 = time.time()
+            runtimes.append(t1 - t0)
+            print(f"Runtime for {c}, n_images={n_images}: {t1-t0}")
+        results[c] = runtimes
+
+        # XLA Mode
+        c = aug.__name__ + " XLA Mode"
+        layer = aug(**aug_args)
+
+        @tf.function(jit_compile=True)
+        def apply_aug(inputs):
+            return layer(inputs)
+
+        runtimes = []
+        print(f"Timing {c}")
+
+        for n_images in num_images:
+            # warmup
+            apply_aug(x_train[:n_images])
+
+            t0 = time.time()
+            r1 = apply_aug(x_train[:n_images])
+            t1 = time.time()
+            runtimes.append(t1 - t0)
+            print(f"Runtime for {c}, n_images={n_images}: {t1-t0}")
+        results[c] = runtimes
+
+    plt.figure()
+    for key in results:
+        plt.plot(num_images, results[key], label=key)
+        plt.xlabel("Number images")
+
+    plt.ylabel("Runtime (seconds)")
+    plt.legend()
+    plt.savefig("comparison.png")
+
+    # So we can actually see more relevant margins
+    del results[aug_candidates[1].__name__]
+    plt.figure()
+    for key in results:
+        plt.plot(num_images, results[key], label=key)
+        plt.xlabel("Number images")
+
+    plt.ylabel("Runtime (seconds)")
+    plt.legend()
+    plt.savefig("comparison_no_old_eager.png")
+
+    # Run unit tests
+    tf.test.main()
diff --git a/keras_cv/layers/preprocessing/channel_shuffle.py b/keras_cv/layers/preprocessing/channel_shuffle.py
index b7ff9eb..0cfed08 100644
--- a/keras_cv/layers/preprocessing/channel_shuffle.py
+++ b/keras_cv/layers/preprocessing/channel_shuffle.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The KerasCV Authors
+# Copyright 2023 The KerasCV Authors
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -14,20 +14,18 @@
 
 import tensorflow as tf
 
-from keras_cv.layers.preprocessing.base_image_augmentation_layer import (
-    BaseImageAugmentationLayer,
+from keras_cv.layers.preprocessing.vectorized_base_image_augmentation_layer import (
+    VectorizedBaseImageAugmentationLayer,
 )
 
 
 @tf.keras.utils.register_keras_serializable(package="keras_cv")
-class ChannelShuffle(BaseImageAugmentationLayer):
+class ChannelShuffle(VectorizedBaseImageAugmentationLayer):
     """Shuffle channels of an input image.
 
     Input shape:
-        The expected images should be [0-255] pixel ranges.
         3D (unbatched) or 4D (batched) tensor with shape:
         `(..., height, width, channels)`, in `"channels_last"` format
-
     Output shape:
         3D (unbatched) or 4D (batched) tensor with shape:
         `(..., height, width, channels)`, in `"channels_last"` format
@@ -36,17 +34,10 @@ class ChannelShuffle(BaseImageAugmentationLayer):
         groups: Number of groups to divide the input channels. Default 3.
         seed: Integer. Used to create a random seed.
 
-    Call arguments:
-        inputs: Tensor representing images of shape
-            `(batch_size, width, height, channels)`, with dtype tf.float32 / tf.uint8,
-            ` or (width, height, channels)`, with dtype tf.float32 / tf.uint8
-        training: A boolean argument that determines whether the call should be run
-            in inference mode or training mode. Default: True.
-
     Usage:
     ```python
     (images, labels), _ = tf.keras.datasets.cifar10.load_data()
-    channel_shuffle = keras_cv.layers.ChannelShuffle()
+    channel_shuffle = ChannelShuffle(groups=3)
     augmented_images = channel_shuffle(images)
     ```
     """
@@ -56,10 +47,40 @@ class ChannelShuffle(BaseImageAugmentationLayer):
         self.groups = groups
         self.seed = seed
 
-    def augment_image(self, image, transformation=None, **kwargs):
-        shape = tf.shape(image)
-        height, width = shape[0], shape[1]
-        num_channels = image.shape[2]
+    def get_random_transformation_batch(self, batch_size, **kwargs):
+        # get batched shuffled indices
+        # for example: batch_size=2; self.group=5
+        # indices = [
+        #     [0, 2, 3, 4, 1],
+        #     [4, 1, 0, 2, 3]
+        # ]
+        indices_distribution = self._random_generator.random_uniform(
+            (batch_size, self.groups)
+        )
+        indices = tf.argsort(indices_distribution, axis=-1)
+        return indices
+
+    def augment_ragged_image(self, image, transformation, **kwargs):
+        # self.augment_images must have
+        # 4D images (batch_size, height, width, channel)
+        # 2D transformations (batch_size, groups)
+        image = tf.expand_dims(image, axis=0)
+        transformation = tf.expand_dims(transformation, axis=0)
+        image = self.augment_images(
+            images=image, transformations=transformation, **kwargs
+        )
+        return tf.squeeze(image, axis=0)
+
+    def augment_images(self, images, transformations, **kwargs):
+        batch_size = tf.shape(images)[0]
+        height, width = images.shape[1], images.shape[2]
+        num_channels = images.shape[3]
+        indices = transformations
+
+        # append batch indexes next to shuffled indices
+        batch_indexs = tf.repeat(tf.range(batch_size), self.groups)
+        batch_indexs = tf.reshape(batch_indexs, (batch_size, self.groups))
+        indices = tf.stack([batch_indexs, indices], axis=-1)
 
         if not num_channels % self.groups == 0:
             raise ValueError(
@@ -69,31 +90,36 @@ class ChannelShuffle(BaseImageAugmentationLayer):
             )
 
         channels_per_group = num_channels // self.groups
-        image = tf.reshape(
-            image, [height, width, self.groups, channels_per_group]
-        )
-        image = tf.transpose(image, perm=[2, 0, 1, 3])
-        image = tf.random.shuffle(image, seed=self.seed)
-        image = tf.transpose(image, perm=[1, 2, 3, 0])
-        image = tf.reshape(image, [height, width, num_channels])
 
-        return image
+        images = tf.reshape(
+            images, [batch_size, height, width, self.groups, channels_per_group]
+        )
+        images = tf.transpose(images, perm=[0, 3, 1, 2, 4])
+        images = tf.gather_nd(images, indices=indices)
+        images = tf.transpose(images, perm=[0, 2, 3, 4, 1])
+        images = tf.reshape(images, [batch_size, height, width, num_channels])
 
-    def augment_bounding_boxes(self, bounding_boxes, **kwargs):
-        return bounding_boxes
+        return images
 
-    def augment_label(self, label, transformation=None, **kwargs):
-        return label
+    def augment_labels(self, labels, transformations, **kwargs):
+        return labels
 
-    def augment_segmentation_mask(
-        self, segmentation_mask, transformation, **kwargs
+    def augment_segmentation_masks(
+        self, segmentation_masks, transformations, **kwargs
     ):
-        return segmentation_mask
+        return segmentation_masks
 
-    def get_config(self):
-        config = super().get_config()
-        config.update({"groups": self.groups, "seed": self.seed})
-        return config
+    def augment_bounding_boxes(self, bounding_boxes, transformations, **kwargs):
+        return bounding_boxes
 
-    def compute_output_shape(self, input_shape):
-        return input_shape
+    def get_config(self):
+        config = {
+            "groups": self.groups,
+            "seed": self.seed,
+        }
+        base_config = super().get_config()
+        return dict(list(base_config.items()) + list(config.items()))
+
+    @classmethod
+    def from_config(cls, config):
+        return cls(**config)
diff --git a/keras_cv/models/convmixer.py b/keras_cv/models/convmixer.py
index ed5fb5a..e5d084b 100644
--- a/keras_cv/models/convmixer.py
+++ b/keras_cv/models/convmixer.py
@@ -59,25 +59,28 @@ MODEL_CONFIGS = {
     },
 }
 
-
 BASE_DOCSTRING = """Instantiates the {name} architecture.
+
     Reference:
         - [Patches Are All You Need?](https://arxiv.org/abs/2201.09792)
-    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(1/255.0)` layer.
-        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
+        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.
         weights: one of `None` (random initialization), a pretrained weight file
             path, or a reference to pre-trained weights (e.g. 'imagenet/classification')
             (see available pre-trained weights in weights.py)
         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`.
@@ -87,87 +90,71 @@ 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.
-        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.
 """
 
 
-def CovnMixer_Layer(dim, kernel_size):
-    """CovnMixer Layer module.
+def apply_conv_mixer_layer(x, dim, kernel_size):
+    """ConvMixerLayer module.
     Args:
-        inputs: Input tensor.
+        x: input tensor.
         dim: integer, filters of the layer in a block.
-        kernel_size: integer, kernel size of the Conv2d layers.
+        kernel_size: integer, kernel size of the Conv2D layers.
     Returns:
-        Output tensor for the CovnMixer Layer.
+        the updated input tensor.
     """
 
-    def apply(x):
-        residual = x
-        x = layers.DepthwiseConv2D(kernel_size=kernel_size, padding="same")(x)
-        x = tf.nn.gelu(x)
-        x = layers.BatchNormalization()(x)
-        x = layers.Add()([x, residual])
-
-        x = layers.Conv2D(dim, kernel_size=1)(x)
-        x = tf.nn.gelu(x)
-        x = layers.BatchNormalization()(x)
-        return x
+    residual = x
+    x = layers.DepthwiseConv2D(kernel_size=kernel_size, padding="same")(x)
+    x = tf.nn.gelu(x)
+    x = layers.BatchNormalization()(x)
+    x = layers.Add()([x, residual])
 
-    return apply
+    x = layers.Conv2D(dim, kernel_size=1)(x)
+    x = tf.nn.gelu(x)
+    x = layers.BatchNormalization()(x)
+    return x
 
 
-def patch_embed(dim, patch_size):
+def apply_patch_embed(x, dim, patch_size):
     """Implementation for Extracting Patch Embeddings.
     Args:
-        inputs: Input tensor.
+        x: input tensor.
+        dim: integer, filters of the layer in a block.
         patch_size: integer, Size of patches.
     Returns:
-        Output tensor for the patch embed.
+        the updated input tensor.
     """
 
-    def apply(x):
-        x = layers.Conv2D(
-            filters=dim, kernel_size=patch_size, strides=patch_size
-        )(x)
-        x = tf.nn.gelu(x)
-        x = layers.BatchNormalization()(x)
-        return x
-
-    return apply
+    x = layers.Conv2D(filters=dim, kernel_size=patch_size, strides=patch_size)(
+        x
+    )
+    x = tf.nn.gelu(x)
+    x = layers.BatchNormalization()(x)
+    return x
 
 
-def ConvMixer(
-    dim,
-    depth,
-    patch_size,
-    kernel_size,
-    include_top,
-    include_rescaling,
-    name="ConvMixer",
-    weights=None,
-    input_shape=(None, None, 3),
-    input_tensor=None,
-    pooling=None,
-    num_classes=None,
-    classifier_activation="softmax",
-):
+@keras.utils.register_keras_serializable(package="keras_cv.models")
+class ConvMixer(keras.Model):
     """Instantiates the ConvMixer architecture.
+
     Args:
-        dim: number of filters.
-        depth: number of CovnMixer Layer.
-        patch_size: Size of the patches.
-        kernel_size: kernel size for conv2d layers.
-        include_rescaling: whether or not to Rescale the inputs. If set to True,
-            inputs will be passed through a `Rescaling(1/255.0)` layer.
-            name: string, model name.
-        include_top: whether to include the fully-connected
+        dim: integer, number of filters.
+        depth: integer, number of ConvMixer Layer.
+        patch_size: integer, size of the patches.
+        kernel_size: integer, kernel size for Conv2D layers.
+        include_top: bool, whether to include the fully-connected
             layer at the top of the network.
-        weights: one of `None` (random initialization),
+        include_rescaling: bool, whether to rescale the inputs. If set to True,
+            inputs will be passed through a `Rescaling(1/255.0)` layer.
+        name: string, optional name to pass to the model, defaults to "ConvMixer".
+        weights: one of `None` (random initialization)
             or the path to the weights file to be loaded.
         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`.
@@ -180,60 +167,107 @@ def ConvMixer(
                 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
-            into, only to be specified if `include_top` is True.
+        num_classes: integer, optional number of classes to classify images
+            into. Only to be specified if `include_top` is True.
         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.
         **kwargs: Pass-through keyword arguments to `tf.keras.Model`.
+
     Returns:
       A `keras.Model` instance.
     """
 
-    if weights and not tf.io.gfile.exists(weights):
-        raise ValueError(
-            "The `weights` argument should be either `None` or the path to the "
-            f"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}. "
-        )
-
-    inputs = utils.parse_model_inputs(input_shape, input_tensor)
-    x = inputs
-
-    if include_rescaling:
-        x = layers.Rescaling(1 / 255.0)(x)
-    x = patch_embed(dim, patch_size)(x)
-
-    for _ in range(depth):
-        x = CovnMixer_Layer(dim, kernel_size)(x)
-
-    if include_top:
-        x = layers.GlobalAveragePooling2D(name="avg_pool")(x)
-        x = layers.Dense(
-            num_classes, activation=classifier_activation, name="predictions"
-        )(x)
-    else:
-        if pooling == "avg":
+    def __init__(
+        self,
+        dim,
+        depth,
+        patch_size,
+        kernel_size,
+        include_top,
+        include_rescaling,
+        name="ConvMixer",
+        weights=None,
+        input_shape=(None, None, 3),
+        input_tensor=None,
+        pooling=None,
+        num_classes=None,
+        classifier_activation="softmax",
+        **kwargs,
+    ):
+        if weights and not tf.io.gfile.exists(weights):
+            raise ValueError(
+                "The `weights` argument should be either `None` or the path to the "
+                f"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 `classes`. "
+                f"Received: 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}. "
+            )
+
+        inputs = utils.parse_model_inputs(input_shape, input_tensor)
+        x = inputs
+
+        if include_rescaling:
+            x = layers.Rescaling(1 / 255.0)(x)
+        x = apply_patch_embed(x, dim, patch_size)
+
+        for _ in range(depth):
+            x = apply_conv_mixer_layer(x, dim, kernel_size)
+
+        if include_top:
             x = layers.GlobalAveragePooling2D(name="avg_pool")(x)
-        elif pooling == "max":
-            x = layers.GlobalMaxPooling2D(name="max_pool")(x)
-
-    model = keras.Model(inputs, x, name=name)
-    if weights is not None:
-        model.load_weights(weights)
-
-    return model
+            x = layers.Dense(
+                num_classes,
+                activation=classifier_activation,
+                name="predictions",
+            )(x)
+        else:
+            if 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.dim = dim
+        self.depth = depth
+        self.patch_size = patch_size
+        self.kernel_size = kernel_size
+        self.include_top = include_top
+        self.include_rescaling = include_rescaling
+        self.input_tensor = input_tensor
+        self.pooling = pooling
+        self.num_classes = num_classes
+        self.classifier_activation = classifier_activation
+
+    def get_config(self):
+        return {
+            "dim": self.dim,
+            "depth": self.depth,
+            "patch_size": self.patch_size,
+            "kernel_size": self.kernel_size,
+            "include_top": self.include_top,
+            "include_rescaling": self.include_rescaling,
+            "name": self.name,
+            "input_shape": self.input_shape[1:],
+            "input_tensor": self.input_tensor,
+            "pooling": self.pooling,
+            "num_classes": self.num_classes,
+            "classifier_activation": self.classifier_activation,
+            "trainable": self.trainable,
+        }
 
 
 def ConvMixer_1536_20(
diff --git a/keras_cv/models/darknet.py b/keras_cv/models/darknet.py
index 7160680..84ee281 100644
--- a/keras_cv/models/darknet.py
+++ b/keras_cv/models/darknet.py
@@ -30,7 +30,7 @@ from keras_cv.models.__internal__.darknet_utils import (
 )
 from keras_cv.models.weights import parse_weights
 
-BASE_DOCSTRING = """Instantiates the {name} architecture.
+BASE_DOCSTRING = """Represents the {name} architecture.
 
     Although the {name} architecture is commonly used for detection tasks, it is
     possible to extract the intermediate dark2 to dark5 layers from the model for
@@ -39,21 +39,22 @@ BASE_DOCSTRING = """Instantiates the {name} architecture.
     Reference:
         - [YoloV3 Paper](https://arxiv.org/abs/1804.02767)
         - [YoloV3 implementation](https://github.com/ultralytics/yolov3)
+
     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(1/255.0)` layer.
-        include_top: whether to include the fully-connected layer at the top of
+        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
+        num_classes: integer, optional number of classes to classify images into. Only to be
             specified if `include_top` is True.
         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,26 +64,17 @@ BASE_DOCSTRING = """Instantiates the {name} architecture.
                 output 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.
-        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.
 """
 
 
-def DarkNet(
-    blocks,
-    include_rescaling,
-    include_top,
-    num_classes=None,
-    weights=None,
-    input_shape=(None, None, 3),
-    input_tensor=None,
-    pooling=None,
-    classifier_activation="softmax",
-    name="DarkNet",
-    **kwargs,
-):
-    """Instantiates the DarkNet architecture.
+@keras.utils.register_keras_serializable(package="keras_cv.models")
+class DarkNet(keras.Model):
+
+    """Represents the DarkNet architecture.
 
     The DarkNet architecture is commonly used for detection tasks. It is
     possible to extract the intermediate dark2 to dark5 layers from the model for
@@ -96,17 +88,17 @@ def DarkNet(
         https://keras.io/guides/transfer_learning/).
 
     Args:
-        blocks: numbers of building blocks from the layer dark2 to layer dark5.
-        include_rescaling: whether or not to Rescale the inputs.If set to True,
+        blocks: integer, numbers of building blocks from the layer dark2 to layer dark5.
+        include_rescaling: bool, whether 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 layer at the top of
+        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 imagesinto, only to be
+        num_classes: integer, optional number of classes to classify images into. Only to be
             specified if `include_top` is True.
-        weights: one of `None` (random initialization), or a pretrained weight
+        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`.
@@ -119,101 +111,145 @@ def DarkNet(
         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.
+        name: string, optional name to pass to the model, defaults to "DarkNet".
 
-        name: (Optional) name to pass to the model.  Defaults to "DarkNet".
     Returns:
         A `keras.Model` instance.
     """
-    if weights and not tf.io.gfile.exists(weights):
-        raise ValueError(
-            "The `weights` argument should be either `None` or the path to the "
-            f"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`. Received: "
-            f"num_classes={num_classes}"
-        )
-
-    inputs = utils.parse_model_inputs(input_shape, input_tensor)
-
-    x = inputs
-    if include_rescaling:
-        x = layers.Rescaling(1 / 255.0)(x)
-
-    # stem
-    x = DarknetConvBlock(
-        filters=32,
-        kernel_size=3,
-        strides=1,
-        activation="leaky_relu",
-        name="stem_conv",
-    )(x)
-    x = ResidualBlocks(filters=64, num_blocks=1, name="stem_residual_block")(x)
-
-    # filters for the ResidualBlock outputs
-    filters = [128, 256, 512, 1024]
-
-    # layer_num is used for naming the residual blocks (starts with dark2, hence 2)
-    layer_num = 2
-
-    for filter, block in zip(filters, blocks):
+
+    def __init__(
+        self,
+        blocks,
+        include_rescaling,
+        include_top,
+        num_classes=None,
+        weights=None,
+        input_shape=(None, None, 3),
+        input_tensor=None,
+        pooling=None,
+        classifier_activation="softmax",
+        name="DarkNet",
+        **kwargs,
+    ):
+        if weights and not tf.io.gfile.exists(weights):
+            raise ValueError(
+                "The `weights` argument should be either `None` or the path to the "
+                f"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`. Received: "
+                f"num_classes={num_classes}"
+            )
+
+        inputs = utils.parse_model_inputs(input_shape, input_tensor)
+
+        x = inputs
+        if include_rescaling:
+            x = layers.Rescaling(1 / 255.0)(x)
+
+        # stem
+        x = DarknetConvBlock(
+            filters=32,
+            kernel_size=3,
+            strides=1,
+            activation="leaky_relu",
+            name="stem_conv",
+        )(x)
         x = ResidualBlocks(
-            filters=filter,
-            num_blocks=block,
-            name=f"dark{layer_num}_residual_block",
+            filters=64, num_blocks=1, name="stem_residual_block"
+        )(x)
+
+        # filters for the ResidualBlock outputs
+        filters = [128, 256, 512, 1024]
+
+        # layer_num is used for naming the residual blocks (starts with dark2, hence 2)
+        layer_num = 2
+
+        for filter, block in zip(filters, blocks):
+            x = ResidualBlocks(
+                filters=filter,
+                num_blocks=block,
+                name=f"dark{layer_num}_residual_block",
+            )(x)
+            layer_num += 1
+
+        # remaining dark5 layers
+        x = DarknetConvBlock(
+            filters=512,
+            kernel_size=1,
+            strides=1,
+            activation="leaky_relu",
+            name="dark5_conv1",
+        )(x)
+        x = DarknetConvBlock(
+            filters=1024,
+            kernel_size=3,
+            strides=1,
+            activation="leaky_relu",
+            name="dark5_conv2",
+        )(x)
+        x = SpatialPyramidPoolingBottleneck(
+            512, activation="leaky_relu", name="dark5_spp"
+        )(x)
+        x = DarknetConvBlock(
+            filters=1024,
+            kernel_size=3,
+            strides=1,
+            activation="leaky_relu",
+            name="dark5_conv3",
         )(x)
-        layer_num += 1
-
-    # remaining dark5 layers
-    x = DarknetConvBlock(
-        filters=512,
-        kernel_size=1,
-        strides=1,
-        activation="leaky_relu",
-        name="dark5_conv1",
-    )(x)
-    x = DarknetConvBlock(
-        filters=1024,
-        kernel_size=3,
-        strides=1,
-        activation="leaky_relu",
-        name="dark5_conv2",
-    )(x)
-    x = SpatialPyramidPoolingBottleneck(
-        512, activation="leaky_relu", name="dark5_spp"
-    )(x)
-    x = DarknetConvBlock(
-        filters=1024,
-        kernel_size=3,
-        strides=1,
-        activation="leaky_relu",
-        name="dark5_conv3",
-    )(x)
-    x = DarknetConvBlock(
-        filters=512,
-        kernel_size=1,
-        strides=1,
-        activation="leaky_relu",
-        name="dark5_conv4",
-    )(x)
-
-    if include_top:
-        x = layers.GlobalAveragePooling2D(name="avg_pool")(x)
-        x = layers.Dense(
-            num_classes, activation=classifier_activation, name="predictions"
+        x = DarknetConvBlock(
+            filters=512,
+            kernel_size=1,
+            strides=1,
+            activation="leaky_relu",
+            name="dark5_conv4",
         )(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)
+        if include_top:
+            x = layers.GlobalAveragePooling2D(name="avg_pool")(x)
+            x = layers.Dense(
+                num_classes,
+                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.blocks = blocks
+        self.include_rescaling = include_rescaling
+        self.include_top = include_top
+        self.num_classes = num_classes
+        self.input_tensor = input_tensor
+        self.pooling = pooling
+        self.classifier_activation = classifier_activation
+
+    def get_config(self):
+        return {
+            "blocks": self.blocks,
+            "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,
+            "classifier_activation": self.classifier_activation,
+            "name": self.name,
+            "trainable": self.trainable,
+        }
 
-    if weights is not None:
-        model.load_weights(weights)
-    return model
+    @classmethod
+    def from_config(cls, config):
+        return cls(**config)
 
 
 def DarkNet21(
diff --git a/keras_cv/models/mlp_mixer.py b/keras_cv/models/mlp_mixer.py
index ca6b89b..5686de0 100644
--- a/keras_cv/models/mlp_mixer.py
+++ b/keras_cv/models/mlp_mixer.py
@@ -50,24 +50,27 @@ MODEL_CONFIGS = {
 }
 
 BASE_DOCSTRING = """Instantiates the {name} architecture.
+
     Reference:
         - [MLP-Mixer: An all-MLP Architecture for Vision](https://arxiv.org/abs/2105.01601)
-    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(1/255.0)` layer.
-        include_top: whether to include the fully-connected layer at the top of the
+        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
+        num_classes: integer, optional number of classes to classify images into. Only to be
             specified if `include_top` is True.
         weights: one of `None` (random initialization), a pretrained weight file
             path, or a reference to pre-trained weights (e.g. 'imagenet/classification')
             (see available pre-trained weights in weights.py)
         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`.
@@ -77,41 +80,41 @@ 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.
-        name: (Optional) name to pass to the model.  Defaults to "{name}".
+        name: string, optional name to pass to the model, defaults to "{name}".
         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.
 """
 
 
-def MLPBlock(mlp_dim, name=None):
+def apply_mlp_block(x, mlp_dim, name=None):
     """An MLP block consisting of two linear layers with GELU activation in
     between.
 
     Args:
+      x: input tensor.
       mlp_dim: integer, the number of units to be present in the first layer.
       name: string, block label.
 
     Returns:
-      a function that takes an input Tensor representing an MLP block.
+      the updated input tensor.
     """
     if name is None:
         name = f"mlp_block_{backend.get_uid('mlp_block')}"
 
-    def apply(x):
-        y = layers.Dense(mlp_dim, name=f"{name}_dense_1")(x)
-        y = layers.Activation("gelu", name=f"{name}_gelu")(y)
-        return layers.Dense(x.shape[-1], name=f"{name}_dense_2")(y)
-
-    return apply
+    y = layers.Dense(mlp_dim, name=f"{name}_dense_1")(x)
+    y = layers.Activation("gelu", name=f"{name}_gelu")(y)
+    return layers.Dense(x.shape[-1], name=f"{name}_dense_2")(y)
 
 
-def MixerBlock(tokens_mlp_dim, channels_mlp_dim, name=None):
+def apply_mixer_block(x, tokens_mlp_dim, channels_mlp_dim, name=None):
     """A mixer block.
 
     Args:
+      x: input tensor.
       tokens_mlp_dim: integer, number of units to be present in the MLP block
         dealing with tokens.
       channels_mlp_dim: integer, number of units to be present in the MLP block
@@ -119,63 +122,45 @@ def MixerBlock(tokens_mlp_dim, channels_mlp_dim, name=None):
       name: string, block label.
 
     Returns:
-      a function that takes an input Tensor representing an MLP block.
+      the updated input tensor.
     """
     if name is None:
         name = f"mixer_block_{backend.get_uid('mlp_block')}"
 
-    def apply(x):
-        y = layers.LayerNormalization()(x)
-        y = layers.Permute((2, 1))(y)
+    y = layers.LayerNormalization()(x)
+    y = layers.Permute((2, 1))(y)
 
-        y = MLPBlock(tokens_mlp_dim, name=f"{name}_token_mixing")(y)
-        y = layers.Permute((2, 1))(y)
-        x = layers.Add()([x, y])
+    y = apply_mlp_block(y, tokens_mlp_dim, name=f"{name}_token_mixing")
+    y = layers.Permute((2, 1))(y)
+    x = layers.Add()([x, y])
 
-        y = layers.LayerNormalization()(x)
-        y = MLPBlock(channels_mlp_dim, name=f"{name}_channel_mixing")(y)
-        return layers.Add()([x, y])
+    y = layers.LayerNormalization()(x)
+    y = apply_mlp_block(y, channels_mlp_dim, name=f"{name}_channel_mixing")
+    return layers.Add()([x, y])
 
-    return apply
 
+@keras.utils.register_keras_serializable(package="keras_cv.models")
+class MLPMixer(keras.Model):
 
-def MLPMixer(
-    input_shape,
-    patch_size,
-    num_blocks,
-    hidden_dim,
-    tokens_mlp_dim,
-    channels_mlp_dim,
-    include_rescaling,
-    include_top,
-    num_classes=None,
-    input_tensor=None,
-    weights=None,
-    pooling=None,
-    classifier_activation="softmax",
-    name=None,
-    **kwargs,
-):
     """Instantiates the MLP Mixer architecture.
 
     Args:
       input_shape: tuple denoting the input shape, (224, 224, 3) for example.
       patch_size: integer denoting the size of the patches to be extracted
         from the inputs (16 for extracting 16x16 patches for example).
-      num_blocks: number of mixer blocks.
-      hidden_dim: dimension to which the patches will be linearly projected.
-      tokens_mlp_dim: dimension of the MLP block responsible for tokens.
-      channels_mlp_dim: dimension of the MLP block responsible for channels.
-      include_rescaling: 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
-        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.
-      weights: one of `None` (random initialization), or a pretrained
+      num_blocks: integer, number of mixer blocks.
+      hidden_dim: integer, dimension to which the patches will be linearly projected.
+      tokens_mlp_dim: integer, dimension of the MLP block responsible for tokens.
+      channels_mlp_dim: integer, dimension of the MLP block responsible for channels.
+      include_rescaling: whether to rescale the inputs. If set to True,
+        inputs will be passed through a `Rescaling(1/255.0)` 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: integer, optional number of classes to classify images
+        into. Only to be specified if `include_top` is True.
+      weights: one of `None` (random initialization) or a pretrained
         weight file path.
-      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`.
@@ -193,80 +178,134 @@ def MLPMixer(
         `classifier_activation=None` to return the logits of the "top" layer.
         When loading pretrained weights, `classifier_activation` can only
         be `None` or `"softmax"`.
-      name: (Optional) name to pass to the model.  Defaults to "DenseNet".
+      name: string, optional name to pass to the model, defaults to "MLPMixer".
 
     Returns:
       A `keras.Model` instance.
     """
-    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 not isinstance(input_shape, tuple):
-        raise ValueError("`input_shape` needs to be tuple.")
-
-    if len(input_shape) != 3:
-        raise ValueError(
-            "`input_shape` needs to contain dimensions for three"
-            " axes: height, width, and channel ((224, 224, 3) for example)."
-        )
-
-    if input_shape[0] != input_shape[1]:
-        raise ValueError("Non-uniform resolutions are not supported.")
-
-    if input_shape[0] % patch_size != 0:
-        raise ValueError(
-            "Input resolution should be divisible by the patch size."
-        )
-
-    inputs = utils.parse_model_inputs(input_shape, input_tensor)
-
-    x = inputs
-    if include_rescaling:
-        x = layers.Rescaling(1 / 255.0)(x)
-
-    x = layers.Conv2D(
-        filters=hidden_dim,
-        kernel_size=(patch_size, patch_size),
-        strides=(patch_size, patch_size),
-        padding="valid",
-        name="patchify_and_projection",
-    )(x)
-    x = layers.Reshape((x.shape[1] * x.shape[2], x.shape[3]))(x)
-
-    for i in range(num_blocks):
-        x = MixerBlock(
-            tokens_mlp_dim, channels_mlp_dim, name=f"mixer_block_{i}"
-        )(x)
 
-    x = layers.LayerNormalization()(x)
-
-    if include_top:
-        x = layers.GlobalAveragePooling1D(name="avg_pool")(x)
-        x = layers.Dense(
-            num_classes, activation=classifier_activation, name="predictions"
+    def __init__(
+        self,
+        input_shape,
+        patch_size,
+        num_blocks,
+        hidden_dim,
+        tokens_mlp_dim,
+        channels_mlp_dim,
+        include_rescaling,
+        include_top,
+        num_classes=None,
+        input_tensor=None,
+        weights=None,
+        pooling=None,
+        classifier_activation="softmax",
+        name="MLPMixer",
+        **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 not isinstance(input_shape, tuple):
+            raise ValueError("`input_shape` needs to be tuple.")
+
+        if len(input_shape) != 3:
+            raise ValueError(
+                "`input_shape` needs to contain dimensions for three"
+                " axes: height, width, and channel ((224, 224, 3) for example)."
+            )
+
+        if input_shape[0] != input_shape[1]:
+            raise ValueError("Non-uniform resolutions are not supported.")
+
+        if input_shape[0] % patch_size != 0:
+            raise ValueError(
+                "Input resolution should be divisible by the patch size."
+            )
+
+        inputs = utils.parse_model_inputs(input_shape, input_tensor)
+
+        x = inputs
+        if include_rescaling:
+            x = layers.Rescaling(1 / 255.0)(x)
+
+        x = layers.Conv2D(
+            filters=hidden_dim,
+            kernel_size=(patch_size, patch_size),
+            strides=(patch_size, patch_size),
+            padding="valid",
+            name="patchify_and_projection",
         )(x)
-
-    elif pooling == "avg":
-        x = layers.GlobalAveragePooling1D(name="avg_pool")(x)
-    elif pooling == "max":
-        x = layers.GlobalMaxPooling1D(name="max_pool")(x)
-
-    model = keras.Model(inputs, x, name=name, **kwargs)
-
-    if weights is not None:
-        model.load_weights(weights)
-    return model
+        x = layers.Reshape((x.shape[1] * x.shape[2], x.shape[3]))(x)
+
+        for i in range(num_blocks):
+            x = apply_mixer_block(
+                x, tokens_mlp_dim, channels_mlp_dim, name=f"mixer_block_{i}"
+            )
+
+        x = layers.LayerNormalization()(x)
+
+        if include_top:
+            x = layers.GlobalAveragePooling1D(name="avg_pool")(x)
+            x = layers.Dense(
+                num_classes,
+                activation=classifier_activation,
+                name="predictions",
+            )(x)
+
+        elif pooling == "avg":
+            x = layers.GlobalAveragePooling1D(name="avg_pool")(x)
+        elif pooling == "max":
+            x = layers.GlobalMaxPooling1D(name="max_pool")(x)
+
+        super().__init__(inputs=inputs, outputs=x, name=name, **kwargs)
+
+        if weights is not None:
+            self.load_weights(weights)
+
+        self.patch_size = patch_size
+        self.num_blocks = num_blocks
+        self.hidden_dim = hidden_dim
+        self.tokens_mlp_dim = tokens_mlp_dim
+        self.channels_mlp_dim = channels_mlp_dim
+        self.include_rescaling = include_rescaling
+        self.include_top = include_top
+        self.num_classes = num_classes
+        self.input_tensor = input_tensor
+        self.pooling = pooling
+        self.classifier_activation = classifier_activation
+
+    def get_config(self):
+        return {
+            "input_shape": self.input_shape[1:],
+            "patch_size": self.patch_size,
+            "num_blocks": self.num_blocks,
+            "hidden_dim": self.hidden_dim,
+            "tokens_mlp_dim": self.tokens_mlp_dim,
+            "channels_mlp_dim": self.channels_mlp_dim,
+            "include_rescaling": self.include_rescaling,
+            "include_top": self.include_top,
+            "num_classes": self.num_classes,
+            "input_tensor": self.input_tensor,
+            "pooling": self.pooling,
+            "classifier_activation": self.classifier_activation,
+            "name": self.name,
+            "trainable": self.trainable,
+        }
+
+    @classmethod
+    def from_config(cls, config):
+        return cls(**config)
 
 
 def MLPMixerB16(
@@ -278,7 +317,7 @@ def MLPMixerB16(
     input_tensor=None,
     weights=None,
     pooling=None,
-    name="mlp_mixer_b16",
+    name="MLPMixerB16",
     **kwargs,
 ):
     """Instantiates the MLPMixerB16 architecture."""
@@ -310,7 +349,7 @@ def MLPMixerB32(
     input_tensor=None,
     weights=None,
     pooling=None,
-    name="mlp_mixer_b32",
+    name="MLPMixerB32",
     **kwargs,
 ):
     """Instantiates the MLPMixerB32 architecture."""
@@ -341,7 +380,7 @@ def MLPMixerL16(
     input_tensor=None,
     weights=None,
     pooling=None,
-    name="mlp_mixer_l16",
+    name="MLPMixerL16",
     **kwargs,
 ):
     """Instantiates the MLPMixerL16 architecture."""
diff --git a/keras_cv/models/regnet.py b/keras_cv/models/regnet.py
index 1dfe773..0d18c02 100644
--- a/keras_cv/models/regnet.py
+++ b/keras_cv/models/regnet.py
@@ -19,8 +19,9 @@ References:
 """
 
 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.layers import SqueezeAndExcite2D
 from keras_cv.models import utils
@@ -203,7 +204,7 @@ MODEL_CONFIGS = {
     },
 }
 
-BASE_DOCSTRING = """Instantiates the {name} architecture.
+BASE_DOCSTRING = """This class represents the {name} architecture.
 
   Reference:
     - [Designing Network Design Spaces](https://arxiv.org/abs/2003.13678)
@@ -255,11 +256,47 @@ BASE_DOCSTRING = """Instantiates the {name} architecture.
 """
 
 
-def Stem(name=None):
+def apply_conv2d_bn(
+    x,
+    filters,
+    kernel_size,
+    strides=1,
+    use_bias=False,
+    groups=1,
+    padding="valid",
+    kernel_initializer="he_normal",
+    batch_norm=True,
+    activation="relu",
+    name="",
+):
+    x = layers.Conv2D(
+        filters,
+        kernel_size,
+        strides=strides,
+        groups=groups,
+        use_bias=use_bias,
+        padding=padding,
+        kernel_initializer=kernel_initializer,
+        name=name,
+    )(x)
+
+    if batch_norm:
+        x = layers.BatchNormalization(
+            momentum=0.9, epsilon=1e-5, name=name + "_bn"
+        )(x)
+
+    if activation is not None:
+        x = layers.Activation(activation, name=name + f"_{activation}")(x)
+
+    return x
+
+
+def apply_stem(x, name=None):
     """Implementation of RegNet stem.
 
     (Common to all model variants)
     Args:
+      x: Tensor, input tensor to the stem
       name: name prefix
 
     Returns:
@@ -268,119 +305,97 @@ def Stem(name=None):
     if name is None:
         name = "stem" + str(backend.get_uid("stem"))
 
-    def apply(x):
-        x = layers.Conv2D(
-            32,
-            (3, 3),
-            strides=2,
-            use_bias=False,
-            padding="same",
-            kernel_initializer="he_normal",
-            name=name + "_stem_conv",
-        )(x)
-        x = layers.BatchNormalization(
-            momentum=0.9, epsilon=1e-5, name=name + "_stem_bn"
-        )(x)
-        x = layers.ReLU(name=name + "_stem_relu")(x)
-        return x
+    x = apply_conv2d_bn(
+        x=x,
+        filters=32,
+        kernel_size=(3, 3),
+        strides=2,
+        padding="same",
+        name=name + "_stem_conv",
+    )
 
-    return apply
+    return x
 
 
-def XBlock(filters_in, filters_out, group_width, stride=1, name=None):
+def apply_x_block(
+    inputs, filters_in, filters_out, group_width, stride=1, name=None
+):
     """Implementation of X Block.
     References:
         - [Designing Network Design Spaces](https://arxiv.org/abs/2003.13678)
 
     Args:
-      filters_in: filters in the input tensor
-      filters_out: filters in the output tensor
-      group_width: group width
-      stride: stride
-      name: name prefix
+      inputs: Tensor, input tensor to the block
+      filters_in: int, filters in the input tensor
+      filters_out: int, filters in the output tensor
+      group_width: int, group width
+      stride: int (or) tuple, stride of Conv layer
+      name: str, name prefix
     Returns:
       Output tensor of the block
     """
     if name is None:
         name = str(backend.get_uid("xblock"))
 
-    def apply(inputs):
-        if filters_in != filters_out and stride == 1:
-            raise ValueError(
-                f"Input filters({filters_in}) and output "
-                f"filters({filters_out}) "
-                f"are not equal for stride {stride}. Input and output filters "
-                f"must be equal for stride={stride}."
-            )
-
-        # Declare layers
-        groups = filters_out // group_width
-
-        if stride != 1:
-            skip = layers.Conv2D(
-                filters_out,
-                (1, 1),
-                strides=stride,
-                use_bias=False,
-                kernel_initializer="he_normal",
-                name=name + "_skip_1x1",
-            )(inputs)
-            skip = layers.BatchNormalization(
-                momentum=0.9, epsilon=1e-5, name=name + "_skip_bn"
-            )(skip)
-        else:
-            skip = inputs
+    if filters_in != filters_out and stride == 1:
+        raise ValueError(
+            f"Input filters({filters_in}) and output "
+            f"filters({filters_out}) "
+            f"are not equal for stride {stride}. Input and output filters "
+            f"must be equal for stride={stride}."
+        )
 
-        # Build block
-        # conv_1x1_1
-        x = layers.Conv2D(
-            filters_out,
-            (1, 1),
-            use_bias=False,
-            kernel_initializer="he_normal",
-            name=name + "_conv_1x1_1",
-        )(inputs)
-        x = layers.BatchNormalization(
-            momentum=0.9, epsilon=1e-5, name=name + "_conv_1x1_1_bn"
-        )(x)
-        x = layers.ReLU(name=name + "_conv_1x1_1_relu")(x)
+    # Declare layers
+    groups = filters_out // group_width
 
-        # conv_3x3
-        x = layers.Conv2D(
-            filters_out,
-            (3, 3),
-            use_bias=False,
+    if stride != 1:
+        skip = apply_conv2d_bn(
+            x=inputs,
+            filters=filters_out,
+            kernel_size=(1, 1),
             strides=stride,
-            groups=groups,
-            padding="same",
-            kernel_initializer="he_normal",
-            name=name + "_conv_3x3",
-        )(x)
-        x = layers.BatchNormalization(
-            momentum=0.9, epsilon=1e-5, name=name + "_conv_3x3_bn"
-        )(x)
-        x = layers.ReLU(name=name + "_conv_3x3_relu")(x)
+            activation=None,
+            name=name + "_skip_1x1",
+        )
+    else:
+        skip = inputs
+
+    # Build block
+    # conv_1x1_1
+    x = apply_conv2d_bn(
+        x=inputs,
+        filters=filters_out,
+        kernel_size=(1, 1),
+        name=name + "_conv_1x1_1",
+    )
 
-        # conv_1x1_2
-        x = layers.Conv2D(
-            filters_out,
-            (1, 1),
-            use_bias=False,
-            kernel_initializer="he_normal",
-            name=name + "_conv_1x1_2",
-        )(x)
-        x = layers.BatchNormalization(
-            momentum=0.9, epsilon=1e-5, name=name + "_conv_1x1_2_bn"
-        )(x)
+    # conv_3x3
+    x = apply_conv2d_bn(
+        x=x,
+        filters=filters_out,
+        kernel_size=(3, 3),
+        strides=stride,
+        groups=groups,
+        padding="same",
+        name=name + "_conv_3x3",
+    )
 
-        x = layers.ReLU(name=name + "_exit_relu")(x + skip)
+    # conv_1x1_2
+    x = apply_conv2d_bn(
+        x=x,
+        filters=filters_out,
+        kernel_size=(1, 1),
+        activation=None,
+        name=name + "_conv_1x1_2",
+    )
 
-        return x
+    x = layers.Activation("relu", name=name + "_exit_relu")(x + skip)
 
-    return apply
+    return x
 
 
-def YBlock(
+def apply_y_block(
+    inputs,
     filters_in,
     filters_out,
     group_width,
@@ -393,99 +408,82 @@ def YBlock(
         - [Designing Network Design Spaces](https://arxiv.org/abs/2003.13678)
 
     Args:
-      filters_in: filters in the input tensor
-      filters_out: filters in the output tensor
-      group_width: group width
-      stride: stride
-      squeeze_excite_ratio: expansion ration for Squeeze and Excite block
-      name: name prefix
+      inputs: Tensor, input tensor to the block
+      filters_in: int, filters in the input tensor
+      filters_out: int, filters in the output tensor
+      group_width: int, group width
+      stride: int (or) tuple, stride of Conv layer
+      squeeze_excite_ratio: float, expansion ratio for Squeeze and Excite block
+      name: str, name prefix
     Returns:
       Output tensor of the block
     """
     if name is None:
         name = str(backend.get_uid("yblock"))
 
-    def apply(inputs):
-        if filters_in != filters_out and stride == 1:
-            raise ValueError(
-                f"Input filters({filters_in}) and output "
-                f"filters({filters_out}) "
-                f"are not equal for stride {stride}. Input and output filters "
-                f"must be equal for stride={stride}."
-            )
-
-        groups = filters_out // group_width
-
-        if stride != 1:
-            skip = layers.Conv2D(
-                filters_out,
-                (1, 1),
-                strides=stride,
-                use_bias=False,
-                kernel_initializer="he_normal",
-                name=name + "_skip_1x1",
-            )(inputs)
-            skip = layers.BatchNormalization(
-                momentum=0.9, epsilon=1e-5, name=name + "_skip_bn"
-            )(skip)
-        else:
-            skip = inputs
+    if filters_in != filters_out and stride == 1:
+        raise ValueError(
+            f"Input filters({filters_in}) and output "
+            f"filters({filters_out}) "
+            f"are not equal for stride {stride}. Input and output filters "
+            f"must be equal for stride={stride}."
+        )
 
-        # Build block
-        # conv_1x1_1
-        x = layers.Conv2D(
-            filters_out,
-            (1, 1),
-            use_bias=False,
-            kernel_initializer="he_normal",
-            name=name + "_conv_1x1_1",
-        )(inputs)
-        x = layers.BatchNormalization(
-            momentum=0.9, epsilon=1e-5, name=name + "_conv_1x1_1_bn"
-        )(x)
-        x = layers.ReLU(name=name + "_conv_1x1_1_relu")(x)
+    groups = filters_out // group_width
 
-        # conv_3x3
-        x = layers.Conv2D(
-            filters_out,
-            (3, 3),
-            use_bias=False,
+    if stride != 1:
+        skip = apply_conv2d_bn(
+            x=inputs,
+            filters=filters_out,
+            kernel_size=(1, 1),
             strides=stride,
-            groups=groups,
-            padding="same",
-            kernel_initializer="he_normal",
-            name=name + "_conv_3x3",
-        )(x)
-        x = layers.BatchNormalization(
-            momentum=0.9, epsilon=1e-5, name=name + "_conv_3x3_bn"
-        )(x)
-        x = layers.ReLU(name=name + "_conv_3x3_relu")(x)
+            activation=None,
+            name=name + "_skip_1x1",
+        )
+    else:
+        skip = inputs
+
+    # Build block
+    # conv_1x1_1
+    x = apply_conv2d_bn(
+        x=inputs,
+        filters=filters_out,
+        kernel_size=(1, 1),
+        name=name + "_conv_1x1_1",
+    )
 
-        # Squeeze-Excitation block
-        x = SqueezeAndExcite2D(
-            filters_out, ratio=squeeze_excite_ratio, name=name
-        )(x)
+    # conv_3x3
+    x = apply_conv2d_bn(
+        x=x,
+        filters=filters_out,
+        kernel_size=(3, 3),
+        strides=stride,
+        groups=groups,
+        padding="same",
+        name=name + "_conv_3x3",
+    )
 
-        # conv_1x1_2
-        x = layers.Conv2D(
-            filters_out,
-            (1, 1),
-            use_bias=False,
-            kernel_initializer="he_normal",
-            name=name + "_conv_1x1_2",
-        )(x)
-        x = layers.BatchNormalization(
-            momentum=0.9, epsilon=1e-5, name=name + "_conv_1x1_2_bn"
-        )(x)
+    # Squeeze-Excitation block
+    x = SqueezeAndExcite2D(filters_out, ratio=squeeze_excite_ratio, name=name)(
+        x
+    )
 
-        x = layers.ReLU(name=name + "_exit_relu")(x + skip)
+    # conv_1x1_2
+    x = apply_conv2d_bn(
+        x=x,
+        filters=filters_out,
+        kernel_size=(1, 1),
+        activation=None,
+        name=name + "_conv_1x1_2",
+    )
 
-        return x
+    x = layers.Activation("relu", name=name + "_exit_relu")(x + skip)
 
-    return apply
+    return x
 
 
-def ZBlock(
+def apply_z_block(
+    inputs,
     filters_in,
     filters_out,
     group_width,
@@ -500,166 +498,159 @@ def ZBlock(
         - [Fast and Accurate Model Scaling](https://arxiv.org/abs/2103.06877).
 
     Args:
-      filters_in: filters in the input tensor
-      filters_out: filters in the output tensor
-      group_width: group width
-      stride: stride
-      squeeze_excite_ratio: expansion ration for Squeeze and Excite block
-      bottleneck_ratio: inverted bottleneck ratio
-      name: name prefix
+      inputs: Tensor, input tensor to the block
+      filters_in: int, filters in the input tensor
+      filters_out: int, filters in the output tensor
+      group_width: int, group width
+      stride: int (or) tuple, stride
+      squeeze_excite_ratio: float, expansion ration for Squeeze and Excite block
+      bottleneck_ratio: float, inverted bottleneck ratio
+      name: str, name prefix
     Returns:
       Output tensor of the block
     """
     if name is None:
         name = str(backend.get_uid("zblock"))
 
-    def apply(inputs):
-        if filters_in != filters_out and stride == 1:
-            raise ValueError(
-                f"Input filters({filters_in}) and output filters({filters_out})"
-                f"are not equal for stride {stride}. Input and output filters "
-                f"must be equal for stride={stride}."
-            )
-
-        groups = filters_out // group_width
+    if filters_in != filters_out and stride == 1:
+        raise ValueError(
+            f"Input filters({filters_in}) and output filters({filters_out})"
+            f"are not equal for stride {stride}. Input and output filters "
+            f"must be equal for stride={stride}."
+        )
 
-        inv_btlneck_filters = int(filters_out / bottleneck_ratio)
+    groups = filters_out // group_width
 
-        # Build block
-        # conv_1x1_1
-        x = layers.Conv2D(
-            inv_btlneck_filters,
-            (1, 1),
-            use_bias=False,
-            kernel_initializer="he_normal",
-            name=name + "_conv_1x1_1",
-        )(inputs)
-        x = layers.BatchNormalization(
-            momentum=0.9, epsilon=1e-5, name=name + "_conv_1x1_1_bn"
-        )(x)
-        x = tf.nn.silu(x)
+    inv_btlneck_filters = int(filters_out / bottleneck_ratio)
 
-        # conv_3x3
-        x = layers.Conv2D(
-            inv_btlneck_filters,
-            (3, 3),
-            use_bias=False,
-            strides=stride,
-            groups=groups,
-            padding="same",
-            kernel_initializer="he_normal",
-            name=name + "_conv_3x3",
-        )(x)
-        x = layers.BatchNormalization(
-            momentum=0.9, epsilon=1e-5, name=name + "_conv_3x3_bn"
-        )(x)
-        x = tf.nn.silu(x)
-
-        # Squeeze-Excitation block
-        x = SqueezeAndExcite2D(
-            inv_btlneck_filters, ratio=squeeze_excite_ratio, name=name
-        )(x)
+    # Build block
+    # conv_1x1_1
+    x = apply_conv2d_bn(
+        x=inputs,
+        filters=inv_btlneck_filters,
+        kernel_size=(1, 1),
+        name=name + "_conv_1x1_1",
+        activation="silu",
+    )
 
-        # conv_1x1_2
-        x = layers.Conv2D(
-            filters_out,
-            (1, 1),
-            use_bias=False,
-            kernel_initializer="he_normal",
-            name=name + "_conv_1x1_2",
-        )(x)
-        x = layers.BatchNormalization(
-            momentum=0.9, epsilon=1e-5, name=name + "_conv_1x1_2_bn"
-        )(x)
+    # conv_3x3
+    x = apply_conv2d_bn(
+        x=x,
+        filters=inv_btlneck_filters,
+        kernel_size=(3, 3),
+        strides=stride,
+        groups=groups,
+        padding="same",
+        name=name + "_conv_3x3",
+        activation="silu",
+    )
 
-        if stride != 1:
-            return x
-        else:
-            return x + inputs
+    # Squeeze-Excitation block
+    x = SqueezeAndExcite2D(
+        inv_btlneck_filters, ratio=squeeze_excite_ratio, name=name
+    )(x)
+
+    # conv_1x1_2
+    x = apply_conv2d_bn(
+        x=x,
+        filters=filters_out,
+        kernel_size=(1, 1),
+        activation=None,
+        name=name + "_conv_1x1_2",
+    )
 
-    return apply
+    if stride != 1:
+        return x
+    else:
+        return x + inputs
 
 
-def Stage(block_type, depth, group_width, filters_in, filters_out, name=None):
+def apply_stage(
+    x, block_type, depth, group_width, filters_in, filters_out, name=None
+):
     """Implementation of Stage in RegNet.
 
     Args:
+      x: Tensor, input tensor to the stage
       block_type: must be one of "X", "Y", "Z"
-      depth: depth of stage, number of blocks to use
-      group_width: group width of all blocks in  this stage
-      filters_in: input filters to this stage
-      filters_out: output filters from this stage
-      name: name prefix
+      depth: int, depth of stage, number of blocks to use
+      group_width: int, group width of all blocks in  this stage
+      filters_in: int, input filters to this stage
+      filters_out: int, output filters from this stage
+      name: str, name prefix
 
     Returns:
-      Output tensor of Stage
+      Output tensor of the block
     """
     if name is None:
         name = str(backend.get_uid("stage"))
 
-    def apply(inputs):
-        x = inputs
-        if block_type == "X":
-            x = XBlock(
-                filters_in,
+    if block_type == "X":
+        x = apply_x_block(
+            x,
+            filters_in,
+            filters_out,
+            group_width,
+            stride=2,
+            name=f"{name}_XBlock_0",
+        )
+        for i in range(1, depth):
+            x = apply_x_block(
+                x,
+                filters_out,
                 filters_out,
                 group_width,
-                stride=2,
-                name=f"{name}_XBlock_0",
-            )(x)
-            for i in range(1, depth):
-                x = XBlock(
-                    filters_out,
-                    filters_out,
-                    group_width,
-                    name=f"{name}_XBlock_{i}",
-                )(x)
-        elif block_type == "Y":
-            x = YBlock(
-                filters_in,
+                name=f"{name}_XBlock_{i}",
+            )
+    elif block_type == "Y":
+        x = apply_y_block(
+            x,
+            filters_in,
+            filters_out,
+            group_width,
+            stride=2,
+            name=name + "_YBlock_0",
+        )
+        for i in range(1, depth):
+            x = apply_y_block(
+                x,
+                filters_out,
                 filters_out,
                 group_width,
-                stride=2,
-                name=name + "_YBlock_0",
-            )(x)
-            for i in range(1, depth):
-                x = YBlock(
-                    filters_out,
-                    filters_out,
-                    group_width,
-                    name=f"{name}_YBlock_{i}",
-                )(x)
-        elif block_type == "Z":
-            x = ZBlock(
-                filters_in,
+                name=f"{name}_YBlock_{i}",
+            )
+    elif block_type == "Z":
+        x = apply_z_block(
+            x,
+            filters_in,
+            filters_out,
+            group_width,
+            stride=2,
+            name=f"{name}_ZBlock_0",
+        )
+        for i in range(1, depth):
+            x = apply_z_block(
+                x,
+                filters_out,
                 filters_out,
                 group_width,
-                stride=2,
-                name=f"{name}_ZBlock_0",
-            )(x)
-            for i in range(1, depth):
-                x = ZBlock(
-                    filters_out,
-                    filters_out,
-                    group_width,
-                    name=f"{name}_ZBlock_{i}",
-                )(x)
-        else:
-            raise NotImplementedError(
-                f"Block type `{block_type}` not recognized."
-                f"block_type must be one of (`X`, `Y`, `Z`). "
+                name=f"{name}_ZBlock_{i}",
             )
-        return x
-
-    return apply
+    else:
+        raise NotImplementedError(
+            f"Block type `{block_type}` not recognized."
+            f"block_type must be one of (`X`, `Y`, `Z`). "
+        )
+    return x
 
 
-def Head(num_classes=None, name=None, activation=None):
+def apply_head(x, num_classes=None, name=None, activation=None):
     """Implementation of classification head of RegNet.
 
     Args:
-      num_classes: number of classes for Dense layer
-      name: name prefix
+      x: Tensor, input to the head block
+      num_classes: int, number of classes for Dense layer
+      name: str, name prefix
 
     Returns:
       Output logits tensor.
@@ -667,55 +658,38 @@ def Head(num_classes=None, name=None, activation=None):
     if name is None:
         name = str(backend.get_uid("head"))
 
-    def apply(x):
-        x = layers.GlobalAveragePooling2D(name=name + "_head_gap")(x)
-        x = layers.Dense(
-            num_classes, name=name + "head_dense", activation=activation
-        )(x)
-        return x
-
-    return apply
-
+    x = layers.GlobalAveragePooling2D(name=name + "_head_gap")(x)
+    x = layers.Dense(
+        num_classes, name=name + "head_dense", activation=activation
+    )(x)
+    return x
 
-def RegNet(
-    depths,
-    widths,
-    group_width,
-    block_type,
-    include_rescaling,
-    include_top,
-    num_classes=None,
-    model_name="regnet",
-    weights=None,
-    input_tensor=None,
-    input_shape=(None, None, 3),
-    pooling=None,
-    classifier_activation="softmax",
-    **kwargs,
-):
-    """Instantiates RegNet architecture given specific configuration.
 
+@keras.utils.register_keras_serializable(package="keras_cv.models")
+class RegNet(tf.keras.Model):
+    """
+    This class represents the architecture of RegNet
     Args:
-        depths: An iterable containing depths for each individual stages.
-        widths: An iterable containing output channel width of each individual
+        depths: iterable, Contains depths for each individual stages.
+        widths: iterable, Contains output channel width of each individual
             stages
-        group_width: Number of channels to be used in each group. See grouped
+        group_width: int, Number of channels to be used in each group. See grouped
             convolutions for more information.
         block_type: Must be one of `{"X", "Y", "Z"}`. For more details see the
             papers "Designing network design spaces" and "Fast and Accurate
             Model Scaling"
-        default_size: Default input image size.
-        model_name: An optional name for the model.
-        include_rescaling: whether or not to Rescale the inputs.If set to True,
+        default_size: tuple (or) list, Default input image size.
+        model_name: str, An optional name for the model.
+        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.
-        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.
-        weights: One of `None` (random initialization), or the path to the
+        weights: str, One of `None` (random initialization), or the path to the
             weights file to be loaded. Defaults to `None`.
-        input_tensor: Optional Keras tensor (i.e. output of `layers.Input()`)
+        input_tensor: Tensor, Optional Keras tensor (i.e. output of `layers.Input()`)
             to use as image input for the model.
         input_shape: Optional shape tuple, defaults to (None, None, 3).
             It should have exactly 3 inputs channels.
@@ -734,76 +708,121 @@ def RegNet(
             use on the "top" layer. Ignored unless `include_top=True`. Set
             `classifier_activation=None` to return the logits of the "top"
             layer. Defaults to `"softmax"`.
-
-    Returns:
-      A `keras.Model` instance.
     """
-    if not (weights is None or tf.io.gfile.exists(weights)):
-        raise ValueError(
-            "The `weights` argument should be either "
-            "`None` (random initialization) "
-            "or the path to the weights file to be loaded."
-        )
 
-    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}. "
-        )
+    def __init__(
+        self,
+        depths,
+        widths,
+        group_width,
+        block_type,
+        include_rescaling,
+        include_top,
+        num_classes=None,
+        model_name="regnet",
+        weights=None,
+        input_tensor=None,
+        input_shape=(None, None, 3),
+        pooling=None,
+        classifier_activation="softmax",
+        **kwargs,
+    ):
+        if not (weights is None or tf.io.gfile.exists(weights)):
+            raise ValueError(
+                "The `weights` argument should be either "
+                "`None` (random initialization) "
+                "or the path to the weights file to be loaded."
+            )
 
-    img_input = utils.parse_model_inputs(input_shape, input_tensor)
-    x = img_input
+        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_rescaling:
-        x = layers.Rescaling(scale=1.0 / 255.0)(x)
-    x = Stem(name=model_name)(x)
+        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}. "
+            )
 
-    in_channels = x.shape[-1]  # Output from Stem
+        img_input = utils.parse_model_inputs(input_shape, input_tensor)
+        x = img_input
 
-    NUM_STAGES = 4
+        if include_rescaling:
+            x = layers.Rescaling(scale=1.0 / 255.0)(x)
+        x = apply_stem(x, name=model_name)
 
-    for stage_index in range(NUM_STAGES):
-        depth = depths[stage_index]
-        out_channels = widths[stage_index]
+        in_channels = x.shape[-1]  # Output from Stem
 
-        x = Stage(
-            block_type,
-            depth,
-            group_width,
-            in_channels,
-            out_channels,
-            name=model_name + "_Stage_" + str(stage_index),
-        )(x)
-        in_channels = out_channels
+        NUM_STAGES = 4
 
-    if include_top:
-        x = Head(num_classes=num_classes, activation=classifier_activation)(x)
-    else:
-        if pooling == "avg":
-            x = layers.GlobalAveragePooling2D()(x)
-        elif pooling == "max":
-            x = layers.GlobalMaxPooling2D()(x)
-
-    model = tf.keras.Model(
-        inputs=img_input, outputs=x, name=model_name, **kwargs
-    )
+        for stage_index in range(NUM_STAGES):
+            depth = depths[stage_index]
+            out_channels = widths[stage_index]
 
-    # Load weights.
-    if weights is not None:
-        model.load_weights(weights)
+            x = apply_stage(
+                x,
+                block_type,
+                depth,
+                group_width,
+                in_channels,
+                out_channels,
+                name=model_name + "_Stage_" + str(stage_index),
+            )
+            in_channels = out_channels
 
-    return model
+        if include_top:
+            x = apply_head(
+                x, num_classes=num_classes, activation=classifier_activation
+            )
+        else:
+            if pooling == "avg":
+                x = layers.GlobalAveragePooling2D()(x)
+            elif pooling == "max":
+                x = layers.GlobalMaxPooling2D()(x)
+
+        super().__init__(inputs=img_input, outputs=x, name=model_name, **kwargs)
+
+        # Load weights.
+        if weights is not None:
+            self.load_weights(weights)
+
+        self.depths = depths
+        self.widths = widths
+        self.group_width = group_width
+        self.block_type = block_type
+        self.include_rescaling = include_rescaling
+        self.include_top = include_top
+        self.num_classes = num_classes
+        self.model_name = model_name
+        self.input_tensor = input_tensor
+        self.pooling = pooling
+        self.classifier_activation = classifier_activation
+
+    def get_config(self):
+        return {
+            "depths": self.depths,
+            "widths": self.widths,
+            "group_width": self.group_width,
+            "block_type": self.block_type,
+            "include_rescaling": self.include_rescaling,
+            "include_top": self.include_top,
+            "num_classes": self.num_classes,
+            "model_name": self.model_name,
+            "input_tensor": self.input_tensor,
+            "input_shape": self.input_shape[1:],
+            "pooling": self.pooling,
+            "classifier_activation": self.classifier_activation,
+            "trainable": self.trainable,
+        }
+
+    @classmethod
+    def from_config(cls, config):
+        return cls(**config)
 
 
 # Instantiating variants
-
-
 def RegNetX002(
     *,
     include_rescaling,
diff --git a/keras_cv/models/segmentation/deeplab.py b/keras_cv/models/segmentation/deeplab.py
index 91a00a9..363075e 100644
--- a/keras_cv/models/segmentation/deeplab.py
+++ b/keras_cv/models/segmentation/deeplab.py
@@ -23,38 +23,26 @@ from keras_cv.models.weights import parse_weights
 
 @keras.utils.register_keras_serializable(package="keras_cv")
 class DeepLabV3(keras.Model):
-    """
-    A segmentation model based on the DeepLab v3.
+    # TODO: add a code example in the docstring.
+    """A segmentation model based on DeepLab v3.
 
     Args:
         num_classes: int, the number of classes for the detection model. Note that
             the num_classes doesn't contain the background class, and the classes
             from the data should be represented by integers with range
             [0, num_classes).
-        backbone: an optional backbone network for the model. Should be a KerasCV model.
-        weights: weights for the complete DeepLabV3 model. one of `None` (random
+        backbone: Optional backbone network for the model. Should be a KerasCV model.
+        weights: Weights for the complete DeepLabV3 model. one of `None` (random
             initialization), a pretrained weight file path, or a reference to
             pre-trained weights (e.g. 'imagenet/classification' or 'voc/segmentation') (see available
             pre-trained weights in weights.py)
-        spatial_pyramid_pooling: also known as Atrous Spatial Pyramid Pooling (ASPP).
+        spatial_pyramid_pooling: Also known as Atrous Spatial Pyramid Pooling (ASPP).
             Performs spatial pooling on different spatial levels in the pyramid, with
             dilation.
-        segmentation_head: an optional `tf.keras.Layer` that predict the segmentation
+        segmentation_head: Optional `keras.Layer` that predict the segmentation
             mask based on feature from backbone and feature from decoder.
     """
 
-    def build(self, input_shape):
-        height = input_shape[1]
-        width = input_shape[2]
-        feature_map_shape = self.backbone.compute_output_shape(input_shape)
-        self.up_layer = tf.keras.layers.UpSampling2D(
-            size=(
-                height // feature_map_shape[1],
-                width // feature_map_shape[2],
-            ),
-            interpolation="bilinear",
-        )
-
     def __init__(
         self,
         num_classes,
@@ -62,15 +50,16 @@ class DeepLabV3(keras.Model):
         spatial_pyramid_pooling=None,
         segmentation_head=None,
         segmentation_head_activation="softmax",
+        weight_decay=0.0001,
         input_shape=(None, None, 3),
         input_tensor=None,
         weights=None,
         **kwargs,
     ):
-        if not isinstance(backbone, tf.keras.layers.Layer):
+        if not isinstance(backbone, keras.layers.Layer):
             raise ValueError(
-                "Backbone need to be a `tf.keras.layers.Layer`, "
-                f"received {backbone}"
+                "Argument `backbone` must be a `keras.layers.Layer` instance. "
+                f"Received instead backbone={backbone} (of type {type(backbone)})."
             )
 
         if weights and not tf.io.gfile.exists(
@@ -89,22 +78,22 @@ class DeepLabV3(keras.Model):
 
         if input_shape[0] is None and input_shape[1] is None:
             raise ValueError(
-                "Input shapes for both the backbone and DeepLabV3 are `None`."
+                "Input shapes for both the backbone and DeepLabV3 cannot be `None`. "
+                f"Received: input_shape={input_shape} and "
+                f"backbone.input_shape={backbone.input_shape[1:]}"
             )
 
         height = input_shape[0]
         width = input_shape[1]
 
-        x = inputs
-
-        feature_map = backbone(x)
+        feature_map = backbone(inputs)
         if spatial_pyramid_pooling is None:
             spatial_pyramid_pooling = SpatialPyramidPooling(
                 dilation_rates=[6, 12, 18]
             )
 
         output = spatial_pyramid_pooling(feature_map)
-        output = tf.keras.layers.UpSampling2D(
+        output = keras.layers.UpSampling2D(
             size=(
                 height // feature_map.shape[1],
                 width // feature_map.shape[2],
@@ -116,7 +105,7 @@ class DeepLabV3(keras.Model):
             segmentation_head = SegmentationHead(
                 num_classes=num_classes,
                 name="segmentation_head",
-                convs=1,
+                convolutions=1,
                 dropout=0.2,
                 kernel_size=1,
                 activation=segmentation_head_activation,
@@ -144,14 +133,22 @@ class DeepLabV3(keras.Model):
         self.spatial_pyramid_pooling = spatial_pyramid_pooling
         self.segmentation_head = segmentation_head
         self.segmentation_head_activation = segmentation_head_activation
-
-    # TODO(tanzhenyu): consolidate how regularization should be applied to KerasCV.
-    def compile(self, weight_decay=0.0001, **kwargs):
         self.weight_decay = weight_decay
-        super().compile(**kwargs)
+
+    def build(self, input_shape):
+        height = input_shape[1]
+        width = input_shape[2]
+        feature_map_shape = self.backbone.compute_output_shape(input_shape)
+        self.up_layer = keras.layers.UpSampling2D(
+            size=(
+                height // feature_map_shape[1],
+                width // feature_map_shape[2],
+            ),
+            interpolation="bilinear",
+        )
 
     def train_step(self, data):
-        images, y_true, sample_weight = tf.keras.utils.unpack_x_y_sample_weight(
+        images, y_true, sample_weight = keras.utils.unpack_x_y_sample_weight(
             data
         )
         with tf.GradientTape() as tape:
@@ -180,10 +177,11 @@ class DeepLabV3(keras.Model):
             "spatial_pyramid_pooling": self.spatial_pyramid_pooling,
             "segmentation_head": self.segmentation_head,
             "segmentation_head_activation": self.segmentation_head_activation,
+            "weight_decay": self.weight_decay,
         }
 
 
-@tf.keras.utils.register_keras_serializable(package="keras_cv")
+@keras.utils.register_keras_serializable(package="keras_cv")
 class SegmentationHead(layers.Layer):
     """Prediction head for the segmentation model
 
@@ -191,28 +189,30 @@ class SegmentationHead(layers.Layer):
     segmentation mask (pixel level classifications) as the output for the model.
 
     Args:
-        num_classes: int, the number of output classes for the prediction. This should
+        num_classes: int, number of output classes for the prediction. This should
             include all the classes (e.g. background) for the model to predict.
-        convs: int, the number of conv2D layers that are stacked before the final
-            classification layer. Default to 2.
-        filters: int, the number of filter/channels for the the conv2D layers. Default
-            to 256.
-        activations: str or 'tf.keras.activations', activation functions between the
-            conv2D layers and the final classification layer. Default to 'relu'
-        output_scale_factor: int, or a pair of ints, for upsample the output mask.
+        convolutions: int, number of `Conv2D` layers that are stacked before the final
+            classification layer. Defaults to 2.
+        filters: int, number of filter/channels for the the conv2D layers.
+            Defaults to 256.
+        activations: str or function, activation functions between the
+            conv2D layers and the final classification layer. Defaults to `"relu"`.
+        output_scale_factor: int, or a pair of ints. Factor for upsampling the output mask.
             This is useful to scale the output mask back to same size as the input
             image. When single int is provided, the mask will be scaled with same
             ratio on both width and height. When a pair of ints are provided, they will
-            be parsed as (height_factor, width_factor). Default to None, which means
+            be parsed as `(height_factor, width_factor)`. Defaults to `None`, which means
             no resize will happen to the output mask tensor.
-        kernel_size: default 3; the kernel_size to be used in each of the `convs` blocks
-        use_bias: default False; whether to use bias or not in each of the `convs` blocks
-                Defaults to none since the blocks use `BatchNormalization` after each conv, rendering
-                bias obsolete
-        activation: default 'softmax', the activation to apply in the classification
-            layer (output of the head)
-
-    Sample code
+        kernel_size: int, the kernel size to be used in each of the convolutional blocks.
+            Defaults to 3.
+        use_bias: boolean, whether to use bias or not in each of the convolutional blocks.
+            Defaults to False since the blocks use `BatchNormalization`
+            after each convolution, rendering bias obsolete.
+        activation: str or function, activation to apply in the classification
+            layer (output of the head). Defaults to `"softmax"`.
+
+    Examples:
+
     ```python
     # Mimic a FPN output dict
     p3 = tf.ones([2, 32, 32, 3])
@@ -230,7 +230,7 @@ class SegmentationHead(layers.Layer):
     def __init__(
         self,
         num_classes,
-        convs=2,
+        convolutions=2,
         filters=256,
         activations="relu",
         dropout=0.0,
@@ -239,24 +239,9 @@ class SegmentationHead(layers.Layer):
         use_bias=False,
         **kwargs,
     ):
-        """
-        Args:
-            num_classes: the number of possible classes for the segmentation map
-            convs: default 2; the number of conv blocks to use in the head (conv2d-batch_norm-activation blocks)
-            filters: default 256; the number of filters in each Conv2D layer
-            activations: default 'relu'; the activation to apply in conv blocks
-            dropout: default 0.0; the dropout to apply between each conv block
-            kernel_size: default 3; the kernel_size to be used in each of the `convs` blocks
-            use_bias: default False; whether to use bias or not in each of the `convs` blocks
-                Defaults to none since the blocks use `BatchNormalization` after each conv, rendering
-                bias obsolete
-            activation: default 'softmax', the activation to apply in the classification
-                layer (output of the head)
-            **kwargs:
-        """
         super().__init__(**kwargs)
         self.num_classes = num_classes
-        self.convs = convs
+        self.convolutions = convolutions
         self.filters = filters
         self.activations = activations
         self.dropout = dropout
@@ -266,10 +251,10 @@ class SegmentationHead(layers.Layer):
 
         self._conv_layers = []
         self._bn_layers = []
-        for i in range(self.convs):
+        for i in range(self.convolutions):
             conv_name = "segmentation_head_conv_{}".format(i)
             self._conv_layers.append(
-                tf.keras.layers.Conv2D(
+                keras.layers.Conv2D(
                     name=conv_name,
                     filters=self.filters,
                     kernel_size=self.kernel_size,
@@ -279,10 +264,10 @@ class SegmentationHead(layers.Layer):
             )
             norm_name = "segmentation_head_norm_{}".format(i)
             self._bn_layers.append(
-                tf.keras.layers.BatchNormalization(name=norm_name)
+                keras.layers.BatchNormalization(name=norm_name)
             )
 
-        self._classification_layer = tf.keras.layers.Conv2D(
+        self._classification_layer = keras.layers.Conv2D(
             name="segmentation_output",
             filters=self.num_classes,
             kernel_size=1,
@@ -294,7 +279,7 @@ class SegmentationHead(layers.Layer):
             dtype=tf.float32,
         )
 
-        self.dropout_layer = tf.keras.layers.Dropout(self.dropout)
+        self.dropout_layer = keras.layers.Dropout(self.dropout)
 
     def call(self, inputs):
         """Forward path for the segmentation head.
@@ -305,7 +290,7 @@ class SegmentationHead(layers.Layer):
         """
         if not isinstance(inputs, dict):
             raise ValueError(
-                f"Expect the inputs to be a dict, but received {inputs}"
+                f"Expect inputs to be a dict. Received instead inputs={inputs}"
             )
 
         lowest_level = next(iter(sorted(inputs)))
@@ -313,17 +298,15 @@ class SegmentationHead(layers.Layer):
         for conv_layer, bn_layer in zip(self._conv_layers, self._bn_layers):
             x = conv_layer(x)
             x = bn_layer(x)
-            x = tf.keras.activations.get(self.activations)(x)
+            x = keras.activations.get(self.activations)(x)
             if self.dropout:
                 x = self.dropout_layer(x)
-
-        x = self._classification_layer(x)
-        return x
+        return self._classification_layer(x)
 
     def get_config(self):
         config = {
             "num_classes": self.num_classes,
-            "convs": self.convs,
+            "convolutions": self.convolutions,
             "filters": self.filters,
             "activations": self.activations,
             "dropout": self.dropout,
diff --git a/pyproject.toml b/pyproject.toml
new file mode 100644
index 0000000..07cfcd6
--- /dev/null
+++ b/pyproject.toml
@@ -0,0 +1,9 @@
+[tool.black]
+line-length = 80
+
+[tool.isort]
+profile = "black"
+force_single_line = "True"
+known_first_party = ["keras_cv", "tests"]
+default_section = "THIRDPARTY"
+line_length = 80
diff --git a/setup.cfg b/setup.cfg
index 8a84d56..d19dd9e 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -35,10 +35,3 @@ ignore =
     W605
     # Ignore for tf.cond lambda
     E731
-
-[isort]
-profile = black
-force_single_line = True
-known_first_party = keras_cv,tests
-default_section = THIRDPARTY
-line_length = 80
diff --git a/shell/format.sh b/shell/format.sh
index ede05d9..c3be8ac 100755
--- a/shell/format.sh
+++ b/shell/format.sh
@@ -1,3 +1,3 @@
 #!/bin/bash
-isort --sl .
-black --line-length 80 .
+isort .
+black .
diff --git a/shell/lint.sh b/shell/lint.sh
index 7386fd6..b51a83b 100755
--- a/shell/lint.sh
+++ b/shell/lint.sh
@@ -31,7 +31,7 @@ then
   exit 1
 fi
 [ $# -eq 0 ] && echo "no issues with flake8"
-black --line-length 80 --check $files
+black --check $files
 if ! [ $? -eq 0 ]
 then
   echo "Please run \"sh shell/format.sh\" to format the code."


Test patch:
diff --git a/keras_cv/layers/preprocessing/channel_shuffle_test.py b/keras_cv/layers/preprocessing/channel_shuffle_test.py
index ad31cba..37064d0 100644
--- a/keras_cv/layers/preprocessing/channel_shuffle_test.py
+++ b/keras_cv/layers/preprocessing/channel_shuffle_test.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The KerasCV Authors
+# Copyright 2023 The KerasCV Authors
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
+import numpy as np
 import tensorflow as tf
 
 from keras_cv.layers.preprocessing.channel_shuffle import ChannelShuffle
@@ -94,3 +94,30 @@ class ChannelShuffleTest(tf.test.TestCase):
         layer = ChannelShuffle(groups=1)
         xs = layer(xs, training=True)
         self.assertTrue(tf.math.reduce_any(xs == 1.0))
+
+    def test_channel_shuffle_on_batched_images_independently(self):
+        image = tf.random.uniform((100, 100, 3))
+        batched_images = tf.stack((image, image), axis=0)
+        layer = ChannelShuffle(groups=3)
+
+        results = layer(batched_images)
+
+        self.assertNotAllClose(results[0], results[1])
+
+    def test_config_with_custom_name(self):
+        layer = ChannelShuffle(name="image_preproc")
+        config = layer.get_config()
+        layer_1 = ChannelShuffle.from_config(config)
+        self.assertEqual(layer_1.name, layer.name)
+
+    def test_output_dtypes(self):
+        inputs = np.array([[[1], [2]], [[3], [4]]], dtype="float64")
+        layer = ChannelShuffle(groups=1)
+        self.assertAllEqual(layer(inputs).dtype, "float32")
+        layer = ChannelShuffle(groups=1, dtype="uint8")
+        self.assertAllEqual(layer(inputs).dtype, "uint8")
+
+    def test_config(self):
+        layer = ChannelShuffle(groups=5)
+        config = layer.get_config()
+        self.assertEqual(config["groups"], 5)
diff --git a/keras_cv/models/segmentation/deeplab_test.py b/keras_cv/models/segmentation/deeplab_test.py
index 5ee511b..2a47e6d 100644
--- a/keras_cv/models/segmentation/deeplab_test.py
+++ b/keras_cv/models/segmentation/deeplab_test.py
@@ -17,6 +17,7 @@ import os
 import pytest
 import tensorflow as tf
 import tensorflow_datasets as tfds
+from tensorflow import keras
 
 from keras_cv.models import ResNet50V2Backbone
 from keras_cv.models import segmentation
@@ -46,7 +47,7 @@ class DeeplabTest(tf.test.TestCase):
     def test_missing_input_shapes(self):
         with self.assertRaisesRegex(
             ValueError,
-            "Input shapes for both the backbone and DeepLabV3 are `None`.",
+            "Input shapes for both the backbone and DeepLabV3 cannot be `None`.",
         ):
             backbone = ResNet50V2Backbone()
             segmentation.DeepLabV3(num_classes=11, backbone=backbone)
@@ -66,7 +67,7 @@ class DeeplabTest(tf.test.TestCase):
         self.assertEquals(output["output"].shape, [1, 64, 64, 11])
 
     def test_mixed_precision(self):
-        tf.keras.mixed_precision.set_global_policy("mixed_float16")
+        keras.mixed_precision.set_global_policy("mixed_float16")
         backbone = ResNet50V2Backbone(
             input_shape=[64, 64, 3],
         )
@@ -78,11 +79,11 @@ class DeeplabTest(tf.test.TestCase):
         output = model(input_image, training=True)
 
         self.assertEquals(output["output"].dtype, tf.float32)
-        tf.keras.mixed_precision.set_global_policy("float32")
+        keras.mixed_precision.set_global_policy("float32")
 
     def test_invalid_backbone_model(self):
         with self.assertRaisesRegex(
-            ValueError, "Backbone need to be a `tf.keras.layers.Layer`"
+            ValueError, "Argument `backbone` must be a `keras.layers.Layer`"
         ):
             segmentation.DeepLabV3(
                 num_classes=11,
@@ -129,10 +130,8 @@ class DeeplabTest(tf.test.TestCase):
         output_res = [96, 96]
         num_images = 11788
 
-        image_resizing = tf.keras.layers.Resizing(
-            target_size[1], target_size[0]
-        )
-        labels_resizing = tf.keras.layers.Resizing(output_res[1], output_res[0])
+        image_resizing = keras.layers.Resizing(target_size[1], target_size[0])
+        labels_resizing = keras.layers.Resizing(output_res[1], output_res[0])
 
         def resize_images_and_masks(data):
             image = tf.image.convert_image_dtype(
@@ -161,7 +160,7 @@ class DeeplabTest(tf.test.TestCase):
         epochs = 1
         model.compile(
             optimizer="adam",
-            loss=tf.keras.losses.BinaryCrossentropy(from_logits=True),
+            loss=keras.losses.BinaryCrossentropy(from_logits=True),
             metrics=["accuracy"],
         )