Repository: quic/aimet (language: python)

Issue:
Keras - Model Preparer Model Subclass Heirarchy



Discussion / hints:

Fix patch:
diff --git a/Docs/api_docs/keras_model_preparer.rst b/Docs/api_docs/keras_model_preparer.rst
index 72b50e7566..3a9b86864e 100644
--- a/Docs/api_docs/keras_model_preparer.rst
+++ b/Docs/api_docs/keras_model_preparer.rst
@@ -121,4 +121,50 @@ The AIMET Keras ModelPreparer API has the following limitations:
         x = x + positions
         return x
 
-  
\ No newline at end of file
+* The AIMET Keras ModelPreparer API may be able to convert models that are inheriting form the Keras Model class or have
+  layers that inherit from the Keras Model class. However, this is not guaranteed. The API will check these layers weights
+  and verify it has the same number of weights as the layers `__init__` defines them. However, if layers defined in the `__init__`
+  are not used in the `call` function, the API will not be able to verify the weights. Furthermore, if a layer defined in the `__init__`
+  is resued, the API will not be able to see both uses. For example, in the ResBlock class below, the `self.relu` is used twice and the
+  API will miss the second use. If the user defines two separate ReLU's, then the API will be able to convert the layer.::
+
+        # Bad Example
+        class ResBlock(tf.keras.Model):
+            def __init__(self, filters, kernel_size):
+                super(ResBlock, self).__init__()
+                self.conv1 = tf.keras.layers.Conv2D(filters, kernel_size, padding='same')
+                self.bn1 = tf.keras.layers.BatchNormalization()
+                self.conv2 = tf.keras.layers.Conv2D(filters, kernel_size, padding='same')
+                self.bn2 = tf.keras.layers.BatchNormalization()
+                self.relu = tf.keras.layers.ReLU()
+    
+            def call(self, input_tensor, training=False):
+                x = self.conv1(input_tensor)
+                x = self.bn1(x, training=training)
+                x = self.relu(x) # First use of self.relu
+                x = self.conv2(x)
+                x = self.bn2(x, training=training)
+                x = self.relu(x) # Second use of self.relu
+                x = tf.keras.layers.add([x, input_tensor])
+                return x
+        
+        # Good Example
+        class ResBlock(tf.keras.Model):
+            def __init__(self, filters, kernel_size):
+                super(ResBlock, self).__init__()
+                self.conv1 = tf.keras.layers.Conv2D(filters, kernel_size, padding='same')
+                self.bn1 = tf.keras.layers.BatchNormalization()
+                self.conv2 = tf.keras.layers.Conv2D(filters, kernel_size, padding='same')
+                self.bn2 = tf.keras.layers.BatchNormalization()
+                self.relu1 = tf.keras.layers.ReLU()
+                self.relu2 = tf.keras.layers.ReLU()
+    
+            def call(self, input_tensor, training=False):
+                x = self.conv1(input_tensor)
+                x = self.bn1(x, training=training)
+                x = self.relu1(x) # First use of self.relu1
+                x = self.conv2(x)
+                x = self.bn2(x, training=training)
+                x = self.relu2(x) # first use of self.relu2
+                x = tf.keras.layers.add([x, input_tensor])
+                return x
\ No newline at end of file
diff --git a/TrainingExtensions/tensorflow/src/python/aimet_tensorflow/keras/batch_norm_fold.py b/TrainingExtensions/tensorflow/src/python/aimet_tensorflow/keras/batch_norm_fold.py
index e4b2af8dc0..67ec251bd6 100644
--- a/TrainingExtensions/tensorflow/src/python/aimet_tensorflow/keras/batch_norm_fold.py
+++ b/TrainingExtensions/tensorflow/src/python/aimet_tensorflow/keras/batch_norm_fold.py
@@ -57,7 +57,7 @@
 
 import aimet_common.libpymo as libpymo
 from aimet_common.utils import AimetLogger
-from aimet_tensorflow.keras.model_preparer import _handle_normal_keras_layer, _update_output_tensors_in_model_layers_connections
+from aimet_tensorflow.keras.model_preparer import _KerasModelPreparer
 from aimet_tensorflow.keras.quant_sim.qc_quantize_wrapper import QcQuantizeWrapper
 from aimet_tensorflow.keras.quant_sim.tensor_quantizer import ParamPerTensorQuantizer
 from aimet_tensorflow.keras.quantsim import QuantizationSimModel
@@ -535,20 +535,21 @@ def wrapped_bn_layer_in_bns_to_remove(layer: tf.keras.layers.Layer) -> bool:
                     # Sort layer_input based on the above dictionary.
                     layer_input = sorted(layer_input, key=lambda current_input, oi=ordered_inputs: oi[current_input.name])
 
-            # Since we are rerouting around the batch normalization layers, we need to temporarily remove the inbound and
-            # outbound nodes of the batch normalization layers so that the model can be built correctly and not duplicate
-            # the non batch normalization layers inbound/outbound nodes.
+            # Since we are rerouting around the batch normalization layers, we need to temporarily remove the inbound
+            # and outbound nodes of the batch normalization layers so that the model can be built correctly and not
+            # duplicate the non batch normalization layers inbound/outbound nodes.
             current_layer._inbound_nodes = []  # pylint: disable=protected-access
-            # Special case for when there is a Lambda opertaion with multiple inputs. For example, z = x + y.
+            # Special case for when there is a Lambda operation with multiple inputs. For example, z = x + y.
             if isinstance(current_layer, TFOpLambda):
-                x = _handle_normal_keras_layer(current_layer, model_layer_connections)
-                current_layer._outbound_nodes = [] # pylint: disable=protected-access
+                kmp = _KerasModelPreparer.get_instance_for_common_layer_passthrough_functions(model_layer_connections)
+                x = kmp._handle_normal_keras_layer(current_layer)  # pylint: disable=protected-access
                 # Updating the Model layer connections
-                _update_output_tensors_in_model_layers_connections(current_layer, x, model, model_layer_connections,
-                                                                   current_layer._outbound_nodes)
+                kmp._update_output_tensors_in_model_layers_connections(  # pylint: disable=protected-access
+                    current_layer, x, model
+                )
             else:
                 x = current_layer(layer_input)
-            current_layer._outbound_nodes = [] # pylint: disable=protected-access
+            current_layer._outbound_nodes = []  # pylint: disable=protected-access
 
             # Set new output tensor (in this case, it will be the same as the original model)
             model_layer_connections[ModelLayerConnectionsProperties.OUTPUT_TENSORS].update({current_layer.name: x})
diff --git a/TrainingExtensions/tensorflow/src/python/aimet_tensorflow/keras/model_preparer.py b/TrainingExtensions/tensorflow/src/python/aimet_tensorflow/keras/model_preparer.py
index 6e8c92ef92..a4505f57c4 100644
--- a/TrainingExtensions/tensorflow/src/python/aimet_tensorflow/keras/model_preparer.py
+++ b/TrainingExtensions/tensorflow/src/python/aimet_tensorflow/keras/model_preparer.py
@@ -37,31 +37,36 @@
 
 """ Implementation to automatically prepare keras models for AIMET by converting them to a functional model """
 import inspect
-from typing import Any, Dict, List, Set, Union
+import logging
+from typing import Any, Dict, List, Set, Union, Optional
 import re
 import numpy as np
 import tensorflow as tf
-import tensorflow.keras.backend as K # pylint: disable=ungrouped-imports
+
+import tensorflow.keras.backend as K
 from packaging import version
 
 if version.parse(tf.version.VERSION) >= version.parse("2.10"):
     # Ignore pylint errors as keras module is not available in TF 2.4
-    from keras.engine.base_layer_utils import is_subclassed # pylint: disable=import-error
-    from keras.engine.functional import Functional # pylint: disable=import-error
-    from keras.engine.keras_tensor import KerasTensor # pylint: disable=import-error
-    from keras.layers.core.tf_op_layer import TFOpLambda # pylint: disable=import-error
-    from keras.layers.merging.base_merge import _Merge as MergeLayersParentClass # pylint: disable=ungrouped-imports, import-error
+    from keras.engine.base_layer_utils import is_subclassed  # pylint: disable=import-error
+    from keras.engine.functional import Functional  # pylint: disable=import-error
+    from keras.engine.keras_tensor import KerasTensor  # pylint: disable=import-error
+    from keras.layers.core.tf_op_layer import TFOpLambda  # pylint: disable=import-error
+    from keras.layers.merging.base_merge import _Merge as MergeLayersParentClass  # pylint: disable=import-error
 else:
     # Ignore pylint errors due to conditional imports
-    from tensorflow.python.keras.engine.base_layer_utils import is_subclassed # pylint: disable=ungrouped-imports
-    from tensorflow.python.keras.engine.keras_tensor import KerasTensor # pylint: disable=ungrouped-imports
-    from tensorflow.python.keras.engine.functional import Functional # pylint: disable=ungrouped-imports
-    from tensorflow.python.keras.layers.core import TFOpLambda # pylint: disable=ungrouped-imports
-    from tensorflow.python.keras.layers.merge import _Merge as MergeLayersParentClass # pylint: disable=ungrouped-imports
+    from tensorflow.python.keras.engine.base_layer_utils import is_subclassed  # pylint: disable=ungrouped-imports
+    from tensorflow.python.keras.engine.keras_tensor import KerasTensor  # pylint: disable=ungrouped-imports
+    from tensorflow.python.keras.engine.functional import Functional  # pylint: disable=ungrouped-imports
+    from tensorflow.python.keras.layers.core import TFOpLambda  # pylint: disable=ungrouped-imports
+    # pylint: disable=ungrouped-imports
+    from tensorflow.python.keras.layers.merge import \
+        _Merge as MergeLayersParentClass
 
 # pylint: disable=wrong-import-position
 from aimet_tensorflow.keras.utils.model_connection_utils import ModelLayerConnections, ModelLayerConnectionsProperties
-from aimet_tensorflow.keras.utils.model_transform_utils import replace_separable_conv_with_depthwise_pointwise, replace_relu6_with_relu
+from aimet_tensorflow.keras.utils.model_transform_utils import replace_separable_conv_with_depthwise_pointwise, \
+    replace_relu6_with_relu
 from aimet_common.utils import AimetLogger
 
 _logger = AimetLogger.get_area_logger(AimetLogger.LogAreas.ModelPreparer)
@@ -69,559 +74,740 @@
 regex_for_camel_case_to_snake_case = re.compile(r'(?<!^)(?=[A-Z])')
 _TEMP_MODEL_NAME = "temp_aimet_intermediate_model"
 
-def _get_original_models_weights_in_functional_model_order(original_model: tf.keras.Model,
-                                                           functional_model: tf.keras.Model,
-                                                           class_names: Set[str]) -> List[np.ndarray]:
-    """
-    Map the original model's weights to the functional model's weights.
-
-    :param original_model: The original model
-    :param functional_model: The functional model
-    :param class_names: The names of the classes that the original model was subclassed from
-    :return: A list of the original model's weights in the order of the functional model's weights
-    """
-
-    # Make the original model's weights into a dictionary for quick lookup by name
-    # The original subclassed layers names are removed to match the new functional model's names
-    original_model_weights = {}
-    for weight in original_model.weights:
-        # pop out class_names of weight name
-        weight_name = weight.name
-        for class_name in class_names:
-            weight_name = weight_name.replace(class_name + '/', '')
-        original_model_weights[weight_name] = weight.numpy()
-
-    # Get the functional model's weights in order as a dictionary for quick lookup where the key is the weight name
-    # and the position of the weight's order is the value
-    functional_model_weight_order = {
-        weight.name: position
-        for position, weight in enumerate(functional_model.weights)
-    }
-
-    # Using the functional model's weights order, get the original model's weights in the same order. The lambda here
-    # uses the weight's name to get position in the functional model's weights order and the sorts the original model's
-    # weights by that position.
-    weights_in_correct_order = [
-        weight for _, weight in
-        sorted(original_model_weights.items(), key=lambda weight_info: functional_model_weight_order[weight_info[0]])
-    ]
-
-    return weights_in_correct_order
-
-
-def _set_functional_models_weights(original_model: tf.keras.Model, functional_model: tf.keras.Model,
-                                   class_names: Set[str]):
-    """
-    Set the functional model's weights to the original model's weights in the correct order
-
-    :param original_model: The original model
-    :param functional_model: The functional model
-    :param class_names: The names of the classes that the original model was subclassed from
-    """
-
-    weights_in_correct_order = \
-        _get_original_models_weights_in_functional_model_order(original_model, functional_model, class_names)
-
-    try:
-        functional_model.set_weights(weights_in_correct_order)
-    except ValueError:
-        _logger.error(
-            "Could not copy weights from original model to functional model. This can occur when "
-            "custom sublayers are defined not in the same order as the sublayers call method. Please ensure that the "
-            "sublayers internal layers are defined in the same order as the sublayers call method.")
-        raise
-
-    _logger.debug("Functional model weights copied.")
-    _logger.info("Model prepared for AIMET in Functional API format.")
-
-
-def _format_input_layer(original_model: tf.keras.Model,
-                        input_layer: Union[tf.keras.layers.InputLayer, List[tf.keras.layers.InputLayer]] = None) \
-                            -> tf.keras.layers.Layer:
-    """
-    This function formats the input layer by either using the original models input layer or the user provided input layer.
-    This function will also raise an error if the model needs a defined input layer to be prepared for AIMET.
-
-    :param original_model: The original model to be copied
-    :param input_layer: The input layer to be used for the functional model
-    :return: The input layer
-    """
-    try:
-        input_layer = original_model.layers[0].input
-
-    except AttributeError:
-        _logger.info("Input layer not found. Using input layer passed in.")
-        if input_layer is None:
-            raise ValueError("The top layer of this model is subclassed. Please provide an input layer via the "
-                             "\'input_layer\' parameter.")
+"""
+This file contains the implementation to automatically prepare keras models for AIMET by converting them to a functional model.
+"""
 
-    return input_layer
 
+class _KerasModelPreparer:
 
-def _get_class_names_in_model(model: Union[tf.keras.Model, tf.keras.layers.Layer]) -> Set[str]:
-    """
-    Helper function to get the class name for a nested layer.
-
-    :param model: the 'layer' or 'model' to get the class name
-    :return: A set containing the class name
-    """
-    return {regex_for_camel_case_to_snake_case.sub("_", model.name).lower()}
-
-
-def _is_nested_layer(layer: tf.keras.layers.Layer) -> bool:
-    """
-    Checks if the given layer is a nested layer.
-
-    :param layer: The layer to check
-    :return: True if the layer is a nested layer, False otherwise
-    """
-    return is_subclassed(layer) or _is_functional_model(layer) or _is_sequential_model(layer)
-
+    def __init__(self, original_model: tf.keras.Model = None, input_layer: tf.keras.layers.InputLayer = None):
+        self.model_outputs = []  # Both normal init and "passthrough" init utilize this
+        if original_model:
+            self.input_layer = self._format_input_layer(original_model, input_layer)
 
-def _is_functional_model(layer: tf.keras.layers.Layer) -> bool:
-    """
-    Checks if the given layer is a functional layer.
-
-    :param layer: The layer to check
-    :return: True if the layer is a functional layer, False otherwise
-    """
-    return isinstance(layer, Functional) and not isinstance(layer, tf.keras.Sequential)
-
-
-def _is_sequential_model(layer: tf.keras.layers.Layer) -> bool:
-    """
-    Checks if the given layer is a sequential layer.
-
-    :param layer: The layer to check
-    :return: True if the layer is a sequential layer, False otherwise
-    """
-    return isinstance(layer, tf.keras.Sequential)
+            if self._inherits_from_keras_model(original_model):
+                _logger.info("This model inherits from tf.keras.Model, connecting model...")
+                self.original_model = self._connect_inherited_model(original_model, input_layer, is_original_model=True)
 
-
-def _set_prepared_models_input_layer(model: tf.keras.Model,
-                                     model_layers_connections: ModelLayerConnectionsProperties.TYPE,
-                                     input_layer: Union[tf.keras.layers.InputLayer, tf.keras.layers.Layer]):
-    """
-    This function sets the input layer of the model to the input layer of the functional model.
-
-    :param model: The original model
-    :param model_layers_connections: The model layers connections dictionary
-    :param input_layer: The input layer used to build input layer for functional model
-    """
-    try:
-        if isinstance(model.input, list):
-            for inp in model.input:
-                model_layers_connections[ModelLayerConnectionsProperties.OUTPUT_TENSORS].update({inp.name: inp})
+            else:
+                self.original_model = original_model
+
+            # Used to fix weight names at end of unwrapping
+            # Originally set to the name of the original model's class in the case that there is an inherited model
+            self.class_names = self._get_class_names_in_model(self.original_model)
+
+            self.model_layers_connections = \
+                ModelLayerConnections.get_model_layers_connection_properties(self.original_model)
+            self._set_prepared_models_input_layer()
+
+            self.original_models_last_layer = self.original_model.layers[-1]
+
+            self.prepared_model = None
+            self.custom_objects = None
+            self.original_weights_in_prepared_model_order = None
+
+    @classmethod
+    def get_instance_for_common_layer_passthrough_functions(
+            cls, model_layers_connections: ModelLayerConnectionsProperties.TYPE
+    ):
+        """
+        Special function to __init__ for classes outside _KerasModelPreparer that want access to useful
+        functions like _handle_normal_keras_layer. ONLY use this for internal use. For normal Keras Model Preparer,
+        please utilize the prepare_model function.
+
+        :param model_layers_connections: Dictionary of Model Layer Connections for the functions to use.
+        :return: A slim instance of _KerasModelPreparer
+        """
+
+        self = cls(original_model=None, input_layer=None)
+        self.model_layers_connections = model_layers_connections
+        return self
+
+    def _get_original_models_weights_in_functional_model_order(self) -> List[np.ndarray]:
+        """
+        Map the original model's weights to the functional model's weights.
+
+        :return: A list of the original model's weights in the order of the functional model's weights
+        """
+        # Make the original model's weights into a dictionary for quick lookup by name
+        # The original subclassed layers names are removed to match the new functional model's names
+        original_model_weights = {}
+        for weight in self.original_model.weights:
+            # pop out class_names of weight name
+            weight_name = weight.name
+            for class_name in self.class_names:
+                weight_name = weight_name.replace(class_name + '/', '')
+            original_model_weights[weight_name] = weight.numpy()
+
+        # Get the functional model's weights in order as a dictionary for quick lookup where the key is the weight name
+        # and the position of the weight's order is the value
+        prepared_model_weight_order = {
+            weight.name: position
+            for position, weight in enumerate(self.prepared_model.weights)
+        }
+
+        # Using the functional model's weights order, get the original model's weights in the same order. The lambda
+        # here uses the weight's name to get position in the functional model's weights order and the sorts the
+        # original model's weights by that position.
+        self.original_weights_in_prepared_model_order = [
+            weight for _, weight in
+            sorted(original_model_weights.items(), key=lambda weight_info: prepared_model_weight_order[weight_info[0]])
+        ]
+
+        return self.original_weights_in_prepared_model_order
+
+    def _set_prepared_models_weights(self):
+        """
+        Set the functional model's weights to the original model's weights in the correct order
+        """
+
+        assert self.prepared_model, "The prepared model must created before setting weights. Please call " \
+                                    "prepare_model() before calling set_weights()."
+
+        try:
+            self.prepared_model.set_weights(self._get_original_models_weights_in_functional_model_order())
+        except ValueError:
+            _logger.error(
+                "Could not copy weights from original model to the prepared model. This can occur when "
+                "custom sublayers are defined not in the same order as the sublayers call method. Please ensure that "
+                "the sublayers internal layers are defined in the same order as the sublayers call method.")
+            raise
+
+        _logger.debug("Functional model weights copied.")
+        _logger.info("Model prepared for AIMET in Functional API format.")
+
+    @staticmethod
+    def _format_input_layer(
+            original_model: tf.keras.Model,
+            input_layer: Union[tf.keras.layers.InputLayer, List[tf.keras.layers.InputLayer]] = None
+    ) -> tf.keras.layers.Layer:
+        """
+        This function formats the input layer by either using the original models input layer or the user provided
+        input layer. This function will also raise an error if the model needs a defined input layer to be prepared
+        for AIMET.
+
+        :param original_model: The original model to be copied
+        :param input_layer: The input layer to be used for the functional model
+        :return: The input layer
+        """
+        if hasattr(original_model, "input"):
+            input_layer = original_model.input
         else:
-            model_layers_connections[ModelLayerConnectionsProperties.OUTPUT_TENSORS].update(
-                {model.input.name: model.input})
-    except AttributeError:
-        # For models that are not connected
-        _logger.info("Model is not connected. Setting input layer to input layer passed in.")
-        model_layers_connections[ModelLayerConnectionsProperties.INBOUND_NODES].update(
-            {model.layers[0].name: [input_layer.name]})
-        model_layers_connections[ModelLayerConnectionsProperties.OUTPUT_TENSORS].update(
-            {input_layer.name: input_layer})
-
-
-def _get_layer_input(layer: tf.keras.layers.Layer,
-                     model_layers_connections: ModelLayerConnectionsProperties.TYPE) -> tf.keras.layers.Layer:
-    """
-    Helper function to get the input layer of a layer.
-
-    :param layer: The layer to get the input layer of
-    :param model_layers_connections: The model layers connections dict
-    :return: The input layer of the layer
-    """
-    try:
-        layer_input = [model_layers_connections[ModelLayerConnectionsProperties.OUTPUT_TENSORS][layer_aux]
-                       for layer_aux in
-                       model_layers_connections[ModelLayerConnectionsProperties.INBOUND_NODES][layer.name]]
-
-        if len(layer_input) == 1:
-            layer_input = layer_input[0]
-    except KeyError:
-        layer_input = _get_most_recently_added_output_tensor(model_layers_connections)
-        model_layers_connections[ModelLayerConnectionsProperties.INBOUND_NODES].update({layer.name: [layer_input.name]})
-        _logger.warning(
-            "Could not find input tensor for layer: %s. Using %s as input, the most recent output tensor.", layer.name,
-            layer_input.name)
-
-    return layer_input
-
-
-def _is_keras_or_tensor_input(arg: Any) -> bool:
-    """
-    Helper function to check if a given argument is a valid Keras tensor.
-
-    :param arg: The argument in question
-    :return: True if it is valid, False if not
-    """
-    if arg is not None:
-        if isinstance(arg, List):
-            return all(isinstance(x, KerasTensor) for x in arg) or all(isinstance(x, tf.Tensor) for x in arg)
-        return isinstance(arg, (KerasTensor, tf.Tensor))
-    return False
-
-
-def _get_updated_call_args(layer: tf.keras.layers.Layer, model_layers_connections: ModelLayerConnectionsProperties.TYPE) -> \
-        List[Union[KerasTensor, List[KerasTensor], Any]]:
-    """
-    Helper function to get the call arguments of a layer.
-
-    :param layer: The layer to get the call arguments of
-    :param model_layers_connections: The model layers connections dict
-    :return: The call arguments of the layer
-    """
-
-    def _is_tf_tensor(arg_in_question: Any) -> bool:
-        return isinstance(arg_in_question, tf.Tensor)
-
-    try:
-        original_call_args = model_layers_connections[ModelLayerConnectionsProperties.CALL_ARGS][layer.name]
-    except KeyError:
-        _logger.warning("Could not find call args for layer: '%s'. Using keras tensor only as input.", layer.name)
-        return [_get_layer_input(layer, model_layers_connections)]
-
-    updated_call_args = []
-    found_keras_tensor = False
-    for arg in original_call_args:
-        if _is_keras_or_tensor_input(arg):
+            _logger.info("Input layer not found. Using input layer passed in.")
+            if input_layer is None:
+                raise ValueError(
+                    "The top layer of this model is subclassed. Please provide an input layer via the "
+                    "\'input_layer\' parameter."
+                )
+
+        if isinstance(input_layer, dict):  # Keras allows passing in tensors via tensor_name : tensor
+            input_layer = [tensor for tensor in input_layer.values()]
+            if len(input_layer) == 1:
+                return input_layer[0]
+
+        return input_layer
+
+    @staticmethod
+    def _get_class_names_in_model(model: Union[tf.keras.Model, tf.keras.layers.Layer]) -> Set[str]:
+        """
+        Helper function to get the class name for a nested layer.
+
+        :param model: the 'layer' or 'model' to get the class name
+        :return: A set containing the class name
+        """
+        return {
+            regex_for_camel_case_to_snake_case.sub("_", name).lower()
+            for name in (model.name, model.__class__.__name__)
+        }
+
+    @staticmethod
+    def _is_nested_layer(layer: tf.keras.layers.Layer) -> bool:
+        """
+        Checks if the given layer is a nested layer.
+
+        :param layer: The layer to check
+        :return: True if the layer is a nested layer, False otherwise
+        """
+        keras_defined_subclassed_layer = is_subclassed(layer)
+        is_aimet_defined_subclassed = keras_defined_subclassed_layer and any(
+            [isinstance(v, tf.keras.layers.Layer) for v in layer.__dict__.values()]
+        )  # check if the subclass is holding subclassed layer attributes. we only care if this is the case.
+
+        return (
+            is_aimet_defined_subclassed or
+            _KerasModelPreparer._is_functional_model(layer) or
+            _KerasModelPreparer._is_sequential_model(layer)
+        )
+
+    @staticmethod
+    def _is_functional_model(layer: tf.keras.layers.Layer) -> bool:
+        """
+        Checks if the given layer is a functional layer.
+
+        :param layer: The layer to check
+        :return: True if the layer is a functional layer, False otherwise
+        """
+        return isinstance(layer, Functional) and not isinstance(layer, tf.keras.Sequential)
+
+    @staticmethod
+    def _is_sequential_model(layer: tf.keras.layers.Layer) -> bool:
+        """
+        Checks if the given layer is a sequential layer.
+
+        :param layer: The layer to check
+        :return: True if the layer is a sequential layer, False otherwise
+        """
+        return isinstance(layer, tf.keras.Sequential)
+
+    def _set_prepared_models_input_layer(self):
+        """
+        This function sets the input layer of the model to the input layer of the functional model.
+        """
+
+        def set_input_layer_factory(input_layer: Union[tf.keras.layers.InputLayer, List[tf.keras.layers.InputLayer]]):
+            if isinstance(input_layer, list):
+                for inp in input_layer:
+                    self.model_layers_connections[ModelLayerConnectionsProperties.OUTPUT_TENSORS].update(
+                        {inp.name: inp}
+                    )
+            else:
+                self.model_layers_connections[ModelLayerConnectionsProperties.OUTPUT_TENSORS].update(
+                    {input_layer.name: input_layer}
+                )
+
+        try:
+            set_input_layer_factory(self.input_layer)
+        except AttributeError:
+            # For models that are not connected
+            _logger.info("Model is not connected. Setting input layer to input layer passed in.")
+
+            input_layer_name = [inp.name for inp in self.input_layer] if isinstance(self.input_layer, list) else \
+                [self.input_layer.name]
+            self.model_layers_connections[ModelLayerConnectionsProperties.INBOUND_NODES].update(
+                {self.original_model.layers[0].name: [*input_layer_name]}
+            )
+
+            set_input_layer_factory(self.input_layer)
+
+    def _get_layer_input(self, layer: tf.keras.layers.Layer) -> tf.keras.layers.Layer:
+        """
+        Helper function to get the input layer of a layer.
+
+        :param layer: The layer to get the input layer of
+        :return: The input layer of the layer
+        """
+        try:
+            layer_input = [
+                self.model_layers_connections[ModelLayerConnectionsProperties.OUTPUT_TENSORS][layer_aux]
+                for layer_aux in
+                self.model_layers_connections[ModelLayerConnectionsProperties.INBOUND_NODES][layer.name]
+            ]
+
+            if len(layer_input) == 1:
+                layer_input = layer_input[0]
+        except KeyError:
+            layer_input = self._get_most_recently_added_output_tensor()
+            self.model_layers_connections[ModelLayerConnectionsProperties.INBOUND_NODES].update(
+                {layer.name: [layer_input.name]}
+            )
+            _logger.warning(
+                "Could not find input tensor for layer: %s. Using %s as input, the most recent output tensor.",
+                layer.name, layer_input.name
+            )
+
+        return layer_input
+
+    @staticmethod
+    def _is_tf_or_keras_tensor_input(arg: Any) -> bool:
+        """
+        Helper function to check if a given argument is a valid Keras tensor.
+
+        :param arg: The argument in question
+        :return: True if it is valid, False if not
+        """
+        if arg is not None:
+            if isinstance(arg, List):
+                return all(isinstance(x, KerasTensor) for x in arg) or all(isinstance(x, tf.Tensor) for x in arg)
+            return isinstance(arg, (KerasTensor, tf.Tensor))
+        return False
+
+    def _get_updated_call_args(self, layer: tf.keras.layers.Layer) -> List[Union[KerasTensor, List[KerasTensor], Any]]:
+        """
+        Helper function to get the call arguments of a layer.
+
+        :param layer: The layer to get the call arguments of
+        :return: The call arguments of the layer
+        """
+
+        def _is_tf_tensor(arg_in_question: Any) -> bool:
+            return isinstance(arg_in_question, tf.Tensor)
+
+        try:
+            original_call_args = self.model_layers_connections[ModelLayerConnectionsProperties.CALL_ARGS][layer.name]
+        except KeyError:
+            _logger.warning("Could not find call args for layer: '%s'. Using keras tensor only as input.", layer.name)
+            return [self._get_layer_input(layer)]
+
+        updated_call_args = []
+        found_keras_tensor = False
+        for arg in original_call_args:
+            if self._is_tf_or_keras_tensor_input(arg):
+
+                if found_keras_tensor and _is_tf_tensor(arg):
+                    updated_call_args.append(arg)
+
+                elif not found_keras_tensor:
+                    layer_input = self._get_layer_input(layer)
+                    if isinstance(layer_input, List):
+                        updated_call_args.extend(layer_input)
+                    else:
+                        updated_call_args.append(layer_input)
+                    found_keras_tensor = True
 
-            if found_keras_tensor and _is_tf_tensor(arg):
+            else:
                 updated_call_args.append(arg)
 
-            elif not found_keras_tensor:
-                layer_input = _get_layer_input(layer, model_layers_connections)
-                if isinstance(layer_input, List):
-                    updated_call_args.extend(layer_input)
+        assert found_keras_tensor, f"No keras tensor found in call args of layer {layer.name}"
+        return updated_call_args
+
+    def _get_call_kwargs(self, layer: tf.keras.layers.Layer) -> Dict[Union[KerasTensor, List[KerasTensor]], Any]:
+        """
+        Helper function to get call keyword arguments for a given layer.
+
+        :param layer: The layer to get the call keyword arguments of
+        :return: The call keyword arguments of the layer
+        """
+        if original_call_kwargs := \
+                self.model_layers_connections[ModelLayerConnectionsProperties.CALL_KWARGS][layer.name]:
+            call_kwargs = {}
+            for key, value in original_call_kwargs.items():
+                # The Keras tensor is already in the call args, so we don't need to add it again. call_kwargs are for
+                # keyword arguments that are not Keras tensors such as 'axis', 'training', etc.
+                if self._is_tf_or_keras_tensor_input(value):
+                    continue
                 else:
-                    updated_call_args.append(layer_input)
-                found_keras_tensor = True
-
+                    call_kwargs[key] = value
         else:
-            updated_call_args.append(arg)
-
-    assert found_keras_tensor, f"No keras tensor found in call args of layer {layer.name}"
-    return updated_call_args
-
-
-def _get_call_kwargs(layer: tf.keras.layers.Layer, model_layers_connections: ModelLayerConnectionsProperties.TYPE) -> \
-        Dict[Union[KerasTensor, List[KerasTensor]], Any]:
-    """
-    Helper function to get call keyword arguments for a given layer.
-
-    :param layer: The layer to get the call keyword arguments of
-    :param model_layers_connections: The model layers connections dict
-    :return: The call keyword arguments of the layer
-    """
-    if original_call_kwargs := model_layers_connections[ModelLayerConnectionsProperties.CALL_KWARGS][layer.name]:
-        call_kwargs = {}
-        for key, value in original_call_kwargs.items():
-            # The Keras tensor is already in the call args, so we don't need to add it again. call_kwargs are for
-            # keyword arguments that are not Keras tensors such as 'axis', 'training', etc.
-            if _is_keras_or_tensor_input(value):
-                continue
+            _logger.debug("No kwargs for layer: '%s'", layer.name)
+            return {}
+        return call_kwargs
+
+    def _update_output_tensors_in_model_layers_connections(
+            self,
+            layer: tf.keras.layers.Layer,
+            new_output_tensor: KerasTensor,
+            model: tf.keras.Model
+    ):
+        """
+        Helper function to update the output tensors in the model layers connections dictionary.
+
+        :param layer: The layer to update the output tensors of
+        :param new_output_tensor: The new output tensor to update with
+        :param model: The model currently being checked. Used to add model outputs
+        """
+        if layer.name != new_output_tensor.name:
+            new_name = new_output_tensor.name
+            old_name_of_inputs = self.model_layers_connections[ModelLayerConnectionsProperties.INBOUND_NODES].pop(
+                layer.name
+            )
+            self.model_layers_connections[ModelLayerConnectionsProperties.INBOUND_NODES].update(
+                {new_name: old_name_of_inputs}
+            )
+
+            # Replace values in model_layers_connections[NetworkDictProperties.INBOUND_NODES] with new_name
+            for value in self.model_layers_connections[ModelLayerConnectionsProperties.INBOUND_NODES].values():
+                if layer.name in value:
+                    idx = value.index(layer.name)
+                    value[idx] = new_name
+
+            self.model_layers_connections[ModelLayerConnectionsProperties.OUTPUT_TENSORS].update(
+                {new_name: new_output_tensor}
+            )
+        else:
+            # Set new output tensor (in this case, it will be the same as the original model)
+            self.model_layers_connections[ModelLayerConnectionsProperties.OUTPUT_TENSORS].update(
+                {layer.name: new_output_tensor}
+            )
+
+        # Save tensor in output list if it is output in the initial model
+        # TODO: Update so that the last conditional is only checked when it's not the last layer.
+        if model.output_names and layer.name in model.output_names:
+            _logger.debug("Layer '%s' added as output layer", layer.name)
+            self.model_outputs.append(new_output_tensor)
+
+    def _get_most_recently_added_output_tensor(self) -> KerasTensor:
+        """
+        Helper function to get the most recently added output tensor from the model layers connections.
+
+        :return: The most recently added output tensor
+        """
+        return next(reversed(self.model_layers_connections[ModelLayerConnectionsProperties.OUTPUT_TENSORS].items()))[-1]
+
+    @staticmethod
+    def _get_temporary_model(layer: tf.keras.layers.Layer, layer_input: tf.keras.layers.Layer) -> tf.keras.Model:
+        """
+        Helper function to create a temporary functional model from a layer.
+
+        :param layer: The layer to create the temporary model from
+        :param layer_input: The input layer of the layer
+        :return: The temporary model
+        """
+
+        def verify_weights(original_layer_weights: Set[tf.Variable], temp_model_weights: Set[tf.Variable]):
+            if missing_weights := original_layer_weights.difference(temp_model_weights):
+                raise ValueError(f"""
+    The number of weights in the temporary model for unwrapping layer '{layer.name}' does not match the
+    number of weights of the original layer. The missing weight(s) are {missing_weights}. This occurs when the Keras 
+    Symbolic tensor passed into the layers call function does not interact with a layer defined inside of the nested 
+    layer. Please refer to the documentation for more information.
+
+    This is the call function that is causing this error:
+    {inspect.getsource(layer.call)}
+    """)
+
+        layer_input = layer_input if isinstance(layer_input, List) else [layer_input]
+        temp_inputs = [
+            tf.keras.layers.Input(shape=inp.shape[1:], name=inp.name.split(':')[0] + "_temp_input")
+            for inp in layer_input
+        ]
+        if len(temp_inputs) == 1:
+            temp_inputs = temp_inputs[0]
+
+        try:
+            if _KerasModelPreparer._inherits_from_keras_model(layer):
+                temp_model = _KerasModelPreparer._connect_inherited_model(layer, temp_inputs)
             else:
-                call_kwargs[key] = value
-    else:
-        _logger.debug("No kwargs for layer: '%s'", layer.name)
-        return {}
-    return call_kwargs
-
-def _update_call_args_in_model_layer_connections(model_layer_connections: ModelLayerConnectionsProperties.TYPE,
-                                                 layer: tf.keras.layers.Layer, new_output_tensor: KerasTensor):
-    """
-    Helper function to update the call args in model layer connections dictionary.
-
-    :param model_layer_connections: The model layers connections dictionary
-    :param layer: The layer to update the output tensors of
-    :param new_output_tensor: The new output tensor to update with
-    """
-    KERAS_SYMBOLIC_TENSORS_INDEX = 0
-    # pylint: disable=protected-access
-    for layer_name, keras_tensor in model_layer_connections[ModelLayerConnectionsProperties.CALL_ARGS].items():
-        keras_tensor = keras_tensor[KERAS_SYMBOLIC_TENSORS_INDEX]
-        if isinstance(keras_tensor, list):
-            for idx, each_keras_tensor in enumerate(keras_tensor):
-                if isinstance(each_keras_tensor, KerasTensor) and \
-                        each_keras_tensor._keras_history.layer.name == layer.name:
-                    model_layer_connections[ModelLayerConnectionsProperties.CALL_ARGS][layer_name]\
-                        [KERAS_SYMBOLIC_TENSORS_INDEX][idx] = new_output_tensor
+                temp_model = tf.keras.Model(inputs=temp_inputs,
+                                            outputs=layer.call(temp_inputs, training=False),
+                                            name=_TEMP_MODEL_NAME)
+            _logger.debug("Model created for layer '%s'", layer.name)
+        except TypeError as e:
+            if "call() got an unexpected keyword argument 'training'" in e.args:
+                _logger.error(
+                    "Model preparer calls subclassed layers call functions with the parameter 'training=False', "
+                    "in the case that the layer behaves differently during evaluation. Please add **kwargs to your "
+                    "call function for layer '%s.'",
+                    layer.name
+                )
+            raise
+
+        temp_model.summary(print_fn=_logger.debug)
+        verify_weights({w.name for w in layer.weights}, {w.name for w in temp_model.weights})
+
+        return temp_model
+
+    @staticmethod
+    def _update_temporary_model_layers_connections_inbound_nodes(
+            temp_model_model_layers_connections: ModelLayerConnectionsProperties.TYPE,
+            temp_model: tf.keras.Model,
+            layer_input: tf.keras.layers.Layer
+    ):
+        """
+        Helper function to update the inbound nodes of the temporary model layers connections dictionary.
+
+        :param temp_model_model_layers_connections: The temporary model layers connections dictionary
+        :param temp_model: The temporary model
+        :param layer_input: The input layer of the layer
+        """
+        temp_model_input_names = [inp.name for inp in temp_model.input] if isinstance(temp_model.input, List) else \
+            [temp_model.input.name]
+        layer_inputs_name = [
+            inp.name for inp in (layer_input if isinstance(layer_input, List) else [layer_input])
+        ]  # pylint: disable=superfluous-parens
+
+        for layers_name, input_tensor_name in temp_model_model_layers_connections[
+                ModelLayerConnectionsProperties.INBOUND_NODES].items():
+            for idx, current_input_name in enumerate(input_tensor_name):
+                if current_input_name in temp_model_input_names:
+                    if len(layer_inputs_name) == 1:  # Special case where the same input is feed in multiple times
+                        temp_model_model_layers_connections[ModelLayerConnectionsProperties.INBOUND_NODES][layers_name][
+                            idx] = layer_inputs_name[0]
+                    else:
+                        temp_model_model_layers_connections[ModelLayerConnectionsProperties.INBOUND_NODES][layers_name][
+                            idx] = layer_inputs_name[idx]
+
+    def _handle_nested_layer(self, layer: tf.keras.layers.Layer) -> KerasTensor:
+        """
+        Helper function to handle nested layers such as subclass, functional, or sequential.
+
+        :param layer: The layer to handle
+        :return: The output tensor of the layer
+        """
+        _logger.debug("Extracting layers for '%s'", layer.name)
+
+        # Converts CamelCase to snake_case of nested layers class name
+        self.class_names.update([layer.name] if self._inherits_from_keras_model(
+            layer) else self._get_class_names_in_model(layer))
+
+        # Create a model based on the nested layer.
+        # This is done with the layer input from the model layers connections dictionary.
+        # 1) The input layer is used to create the temporary functional model
+        # 2) The input layer is used in the nested layers call function as a symbolic tensor to get internal layers
+        layer_input = self._get_layer_input(layer)
+        temp_model = tf.keras.models.clone_model(self._get_temporary_model(layer, layer_input))
+
+        # Get the model layers connections dictionary for the temporary model and merge it with the model layers
+        # connections dictionary for the functional model. This is done, so we can keep track of the sublayer and
+        # their inputs and outputs.
+        temp_model_model_layers_connections = ModelLayerConnections.get_model_layers_connection_properties(temp_model)
+        self._update_temporary_model_layers_connections_inbound_nodes(
+            temp_model_model_layers_connections, temp_model, layer_input
+        )
+
+        self.model_layers_connections = ModelLayerConnections.merge_model_layers_connections(
+            self.model_layers_connections, temp_model_model_layers_connections
+        )
+
+        return self._prepare_model_helper(temp_model)
+
+    def _handle_normal_keras_layer(self, layer: tf.keras.layers.Layer) -> KerasTensor:
+        """
+        Helper function to handle normal keras layers. This function will create a new output tensor for the layer
+        and return it.
+
+        :param layer: The layer to create the output tensor for
+        :return: The output tensor of the layer
+        """
+        call_args = self._get_updated_call_args(layer)
+
+        if isinstance(layer, TFOpLambda):
+            if call_kwargs := self._get_call_kwargs(layer):
+                # Special case for 'tf.concat' that takes a list of inputs with kwargs attached
+                # may need to updated in the future
+
+                if "concat" in layer.name:
+                    new_output_tensor = layer.call([*call_args], **call_kwargs)
+                else:
+                    new_output_tensor = layer.call(*call_args, **call_kwargs)
+            else:
+                new_output_tensor = layer.call(*call_args)
+        # Special case for "Merge" layers that take a list of inputs such as "tf.keras.layers.Concatenate" and
+        # "tf.keras.layers.Add"
+        elif isinstance(layer, MergeLayersParentClass):
+            new_output_tensor = layer(call_args)
         else:
-            if isinstance(keras_tensor, KerasTensor) and keras_tensor._keras_history.layer.name == layer.name:
-                model_layer_connections[ModelLayerConnectionsProperties.CALL_ARGS][layer_name] = (new_output_tensor,)
-
-
-def _update_output_tensors_in_model_layers_connections(layer: tf.keras.layers.Layer, new_output_tensor: KerasTensor,
-                                                       model: tf.keras.Model,
-                                                       model_layers_connections: ModelLayerConnectionsProperties.TYPE,
-                                                       model_outputs: List[KerasTensor]):
-    """
-    Helper function to update the output tensors in the model layers connections dictionary. It also updates
-    the call_args of model layer connections
-
-    :param layer: The layer to update the output tensors of
-    :param new_output_tensor: The new output tensor to update with
-    :param model: The model currently being checked. Used to add model outputs
-    :param model_layers_connections: The model layers connections dictionary
-    :param model_outputs: A list tracking the currently being checked models output tensors
-    """
-    if layer.name != new_output_tensor.name:
-        new_name = new_output_tensor.name
-        old_name_of_inputs = model_layers_connections[ModelLayerConnectionsProperties.INBOUND_NODES].pop(
-            layer.name)
-        model_layers_connections[ModelLayerConnectionsProperties.INBOUND_NODES].update(
-            {new_name: old_name_of_inputs})
-
-        # Replace values in model_layers_connections[NetworkDictProperties.INBOUND_NODES] with new_name
-        for value in model_layers_connections[ModelLayerConnectionsProperties.INBOUND_NODES].values():
-            if layer.name in value:
-                idx = value.index(layer.name)
-                value[idx] = new_name
-
-        model_layers_connections[ModelLayerConnectionsProperties.OUTPUT_TENSORS].update({new_name: new_output_tensor})
-    else:
-        # Set new output tensor (in this case, it will be the same as the original model)
-        model_layers_connections[ModelLayerConnectionsProperties.OUTPUT_TENSORS].update(
-            {layer.name: new_output_tensor})
-
-    # Updating the CALL_ARGS of model layer connections with new output tensor
-    _update_call_args_in_model_layer_connections(model_layers_connections, layer, new_output_tensor)
-
-    # Save tensor in output list if it is output in the initial model
-    # TODO: Update so that the last conditional is only checked when it's not the last layer.
-    if model.output_names and layer.name in model.output_names:
-        _logger.debug("Layer '%s' added as output layer", layer.name)
-        model_outputs.append(new_output_tensor)
-
-
-def _get_most_recently_added_output_tensor(model_layers_connections: ModelLayerConnectionsProperties.TYPE) \
-        -> KerasTensor:
-    """
-    Helper function to get the most recently added output tensor from the model layers connections.
-
-    :param model_layers_connections: The model layers connections dictionary
-    :return: The most recently added output tensor
-    """
-    return next(reversed(model_layers_connections[ModelLayerConnectionsProperties.OUTPUT_TENSORS].items()))[-1]
-
-
-def _get_temporary_model(layer: tf.keras.layers.Layer, layer_input: tf.keras.layers.Layer) -> tf.keras.Model:
-    """
-    Helper function to create a temporary functional model from a layer.
-
-    :param layer: The layer to create the temporary model from
-    :param layer_input: The input layer of the layer
-    :return: The temporary model
-    """
-
-    def verify_weights(original_layer_weights: Set[tf.Variable], temp_model_weights: Set[tf.Variable]):
-        if missing_weights := original_layer_weights.difference(temp_model_weights):
-            raise  ValueError(f"""
-The number of weights in the temporary model for unwrapping layer '{layer.name}' does not match the
-number of weights of the original layer. The missing weight(s) are {missing_weights}. This occurs when the Keras Symbolic tensor
-passed into the layers call function does not interact with a layer defined inside of the nested layer. Please refer to
-the documentation for more information.
-
-This is the call function that is causing this error:
-{inspect.getsource(layer.call)}
-""")
-
-    temp_input = tf.keras.layers.Input(shape=layer_input.shape[1:], name=layer_input.name + "_temp_input")
-    temp_model = tf.keras.Model(inputs=[temp_input],
-                                outputs=layer.call(temp_input, training=False),
-                                name=_TEMP_MODEL_NAME)
-    _logger.debug("Model created for layer '%s'", layer.name)
-    temp_model.summary(print_fn=_logger.debug)
-
-    verify_weights({w.name for w in layer.weights}, {w.name for w in temp_model.weights})
-
-    return temp_model
-
-
-def _update_temporary_model_layers_connections_inbound_nodes(
-        temp_model_model_layers_connections: ModelLayerConnectionsProperties.TYPE, temp_model: tf.keras.Model,
-        layer_input: tf.keras.layers.Layer):
-    """
-    Helper function to update the inbound nodes of the temporary model layers connections dictionary.
-
-    :param temp_model_model_layers_connections: The temporary model layers connections dictionary
-    :param temp_model: The temporary model
-    :param layer_input: The input layer of the layer
-    """
-    for layers_name, input_tensor_name in temp_model_model_layers_connections[ModelLayerConnectionsProperties.INBOUND_NODES].items():
-        for idx, current_input_name in enumerate(input_tensor_name):
-            if current_input_name == temp_model.input.name:
-                temp_model_model_layers_connections[ModelLayerConnectionsProperties.INBOUND_NODES][layers_name][
-                    idx] = layer_input.name
-
-
-def _handle_nested_layer(layer: tf.keras.layers.Layer, model_layers_connections: Dict, class_names: Set[str],
-                         model_outputs: List[KerasTensor], original_models_last_layer: tf.keras.layers.Layer) -> KerasTensor:
-    """
-    Helper function to handle nested layers such as subclass, functional, or sequential.
-
-    :param layer: The layer to handle
-    :param model_layers_connections: The model layers connections dictionary
-    :param class_names: The list of class names
-    :param model_outputs: The list of model outputs
-    :return: The output tensor of the layer
-    """
-    _logger.debug("Subclass layer '%s' found. Extracting layers.", layer.name)
-    # Converts CamelCase to snake_case of nested layers class name
-    class_names.update(_get_class_names_in_model(layer))
-
-    # Create a model based on the nested layer.
-    # This is done with the layer input from the model layers connections dictionary.
-    # 1) The input layer is used to create the temporary functional model
-    # 2) The input layer is used in the nested layers call function as a symbolic tensor to get internal layers
-    layer_input = _get_layer_input(layer, model_layers_connections)
-    temp_model = tf.keras.models.clone_model(_get_temporary_model(layer, layer_input))
-
-    # Get the model layers connections dictionary for the temporary model and merge it with the model layers connections dictionary for the
-    # functional model. This is done, so we can keep track of the sublayer and their inputs and outputs.
-    temp_model_model_layers_connections = ModelLayerConnections.get_model_layers_connection_properties(temp_model)
-    _update_temporary_model_layers_connections_inbound_nodes(temp_model_model_layers_connections,
-                                                             temp_model, layer_input)
-
-    model_layers_connections = ModelLayerConnections.merge_model_layers_connections(model_layers_connections,
-                                                                                    temp_model_model_layers_connections)
-
-    return _prepare_model_helper(temp_model,
-                                 class_names,
-                                 model_layers_connections,
-                                 model_outputs,
-                                 original_models_last_layer)
-
-
-def _handle_normal_keras_layer(layer: tf.keras.layers.Layer,
-                               model_layers_connections: ModelLayerConnectionsProperties.TYPE) -> KerasTensor:
-    """
-    Helper function to handle normal keras layers. This function will create a new output tensor for the layer
-    and return it.
+            new_output_tensor = layer(*call_args)
+
+        return new_output_tensor
+
+    def _prepare_model_helper(self, model: tf.keras.Model) -> KerasTensor:
+        """
+        Helper function to recursively prepare a model. This function will be recursively called if a nested layer is
+        found. This function will extract the layers from the nested layer and add them to the functional model.
+        Otherwise, it will add the layer to the functional model.
+
+        :param model: The model to prepare
+        :return: The last layer of the model
+        """
+        for current_layer in model.layers:
+            _logger.debug("Processing layer: '%s'", current_layer.name)
+            # Skip input layers
+            if isinstance(current_layer, tf.keras.layers.InputLayer):
+                continue
 
-    :param layer: The layer to create the output tensor for
-    :param model_layers_connections: The models layer connections dictionary
-    :return: The output tensor of the layer
-    """
-    call_args = _get_updated_call_args(layer, model_layers_connections)
-
-    if isinstance(layer, TFOpLambda):
-        if call_kwargs := _get_call_kwargs(layer, model_layers_connections):
-            # Special case for 'tf.concat' that takes a list of inputs with kwargs attached
-            # may need to updated in the future...
-            if 'concat' in layer.name:
-                new_output_tensor = layer.call([*call_args], **call_kwargs)
+            # If the current layer is either a subclassed layer, functional model or sequential model, we need to
+            # extract the layers from the nested layer and add them to the functional model.
+            if self._is_nested_layer(current_layer):
+                new_output_tensor = self._handle_nested_layer(current_layer)
+                # If we are at the end of the original model, we want the model_outputs to be the end model outputs
+                if current_layer == self.original_models_last_layer:
+                    _logger.debug(
+                        "Last layer was a nested layer. "
+                        "Using temp model's output from _handle_nested_layer as model_output"
+                    )
+                    continue
+                self.model_outputs.clear()
             else:
-                new_output_tensor = layer.call(*call_args, **call_kwargs)
+                new_output_tensor = self._handle_normal_keras_layer(current_layer)
+
+            self._update_output_tensors_in_model_layers_connections(current_layer, new_output_tensor, model)
+        return new_output_tensor
+
+    def prepare_model(self):
+        """
+        Function to get the prepared model. This function sets up the input layer and calls the helper function to
+        recursively prepare the model.
+        """
+        _ = self._prepare_model_helper(self.original_model)
+
+        # If the model outputs are empty, then we need to get the most recently added output tensor. This is the case
+        # when a model might be sparse and not fully connected or when a Functional model is inside an inherited model.
+        if not self.model_outputs:
+            _logger.warning(
+                "No model outputs found. This usually occurs when a models is made by inheriting from "
+                "'tf.keras.Model' and placing a Functional model inside. Using most recently added output tensor as "
+                "prepared models output."
+            )
+            self.model_outputs = self._get_most_recently_added_output_tensor()
+
+        setattr(self, "prepared_model", tf.keras.Model(
+            inputs=self.input_layer,
+            outputs=self.model_outputs,
+            name=f"{self.original_model.name}_prepared"
+        ))
+
+        # Cloning model to remove any references to the original model
+        K.clear_session()  # To avoid name conflicts
+        self.prepared_model = tf.keras.models.clone_model(self.prepared_model)
+        setattr(
+            self, "custom_objects",  # For acceptable subclass layers
+            self._get_models_custom_objects(self.prepared_model)
+        )
+        _logger.info("Prepared Model Summary: \n")
+        self.prepared_model.summary(print_fn=_logger.info)
+
+        # Copying over weights from original model to functional model
+        _logger.debug("Final class_names: %s", self.class_names)
+        self._set_prepared_models_weights()
+
+        # Extra prepare step to replace Separable Conv's with Depthwise Pointwise pattern.
+        self.prepared_model, _ = replace_separable_conv_with_depthwise_pointwise(
+            self.prepared_model,
+            custom_objects=self.custom_objects
+        )
+        self.prepared_model, _ = replace_relu6_with_relu(
+            self.prepared_model,
+            custom_objects=self.custom_objects
+        )
+
+        self.verify_prepared_model()
+
+    @staticmethod
+    def _get_models_custom_objects(model: tf.keras.Model) -> Optional[Dict[str, tf.keras.layers.Layer]]:
+        """
+        Helper function to return a models `custom_objects` if there are any present in the model.
+
+        :param model: The model to check
+        :return: A dictionary {layer name : layer obj} of the custom objects or None if there are not any
+        """
+
+        return {
+            layer.__class__.__name__: layer.__class__
+            for layer in model.layers
+            if not getattr(layer, "__module__", None).split(".")[0] == "keras" and  # TF 2.10.1 and up
+            not getattr(layer, "__module__", None).split(".")[0] == "tensorflow"    # TF 2.4.3 support
+        } or None
+
+    @staticmethod
+    def _model_has_nested_layers(model: tf.keras.Model) -> bool:
+        """
+        Helper function to check if a model is needed to be prepared or not based on if the model has nested layers such as
+        subclass, functional, or sequential.
+
+        :param model: The model to check
+        :return: If the model needs to be prepared or not
+        """
+        for layer in model.layers:
+            if _KerasModelPreparer._is_nested_layer(layer):
+                return True
+        return False
+
+    @staticmethod
+    def _inherits_from_keras_model(model: tf.keras.Model) -> bool:
+        """
+        Helper function to check if a model itself is inheriting from tf.keras.Model. If so, then the model needs to connected.
+
+        :param model: The model to check.
+        :return: If the model is inheriting from tf.keras.Model
+        """
+
+        return (
+            type(model).__bases__[0] == tf.keras.Model and
+            not _KerasModelPreparer._is_functional_model(model) and
+            not _KerasModelPreparer._is_sequential_model(model)
+        )
+
+    @staticmethod
+    def _connect_inherited_model(model: tf.keras.Model, input_layer: Union[
+            tf.keras.layers.InputLayer, List[tf.keras.layers.InputLayer]],
+                                 is_original_model: bool = False) -> tf.keras.Model:
+        """
+        Function to loop through models that inherit from tf.keras.Model and therefore could potentially have no
+        outbound nodes.
+
+        :param model: Model to connect.
+        :param input_layer: The input layer to connect the model.
+        :param is_original_model: Flag to clone the model if the original model is the one passed in.
+        This is to fix naming issues. Otherwise, the model is not cloned.
+        :return: A model with the outbound nodes generated.
+        """
+
+        # TODO: Fix case where the layers are all the same. Maybe user has to?
+        model = tf.keras.Model(inputs=input_layer, outputs=model.call(input_layer), name=_TEMP_MODEL_NAME)
+        if is_original_model:
+            try:
+                return tf.keras.models.clone_model(model)
+            except TypeError as e:
+                _logger.error("The layer %s inherits from tf.keras.Model and has layer that does not have a "
+                              "`get_config` defined. Due to this, Keras cannot clone this layer. Please override the "
+                              "`get_config` function and provide the missing keys mentioned in the Keras error logs.",
+                              model.name)
+                raise e
+        return model
+
+    def verify_prepared_model(self):
+        """
+        Function to verify that the prepared model is correct. This function will check that the prepared model has
+        the same weights as the original model and that the prepared model has the same outputs as the original model.
+        """
+
+        # Check that the prepared model has the same number of parameters as the original model
+        assert self.prepared_model.count_params() == self.original_model.count_params(), \
+            "Prepared model and original model do not have the same number of parameters"
+        _logger.debug("Prepared model and original model have the same number of parameters")
+
+        # Check the weights of the prepared model and the original model
+        for original_weight, prepared_weight in zip(
+                self.original_weights_in_prepared_model_order, self.prepared_model.get_weights()):
+            np.testing.assert_array_equal(
+                original_weight, prepared_weight,
+                err_msg="Weights of prepared model and original model do not match"
+            )
+        _logger.debug("Weights of prepared model and original model match")
+
+        # Create a random input to test the prepared model
+        if isinstance(self.prepared_model.input_shape, List):
+            random_input = []
+            for current_input_shape in self.original_model.input_shape:
+                input_shape = [shape if shape is not None else 1 for shape in current_input_shape]
+                random_input.append(np.random.rand(*input_shape).astype(np.float32))
         else:
-            new_output_tensor = layer.call(*call_args)
-    # Special case for "Merge" layers that take a list of inputs such as "tf.keras.layers.Concatenate" and "tf.keras.layers.Add"
-    elif isinstance(layer, MergeLayersParentClass):
-        new_output_tensor = layer(call_args)
-    else:
-        new_output_tensor = layer(*call_args)
-
-    return new_output_tensor
-
-
-def _prepare_model_helper(model: tf.keras.Model, class_names: Set[str],
-                          model_layers_connections: ModelLayerConnectionsProperties.TYPE,
-                          model_outputs: List[KerasTensor],
-                          original_models_last_layer: tf.keras.layers.Layer) -> KerasTensor:
-    """
-    Helper function to recursively prepare a model. This function will be recursively called if a nested layer is
-    found. This function will extract the layers from the nested layer and add them to the functional model.
-    Otherwise, it will add the layer to the functional model.
-
-    :param model: The model to prepare
-    :param class_names: The names of the classes that the original model was subclassed from
-    :param model_layers_connections: The model layers connections dict
-    :param model_outputs: The list tracking the models output tensors
-    :return: The last layer of the model
-    """
-    for current_layer in model.layers:
-        _logger.debug("Processing layer: '%s'", current_layer.name)
-        # Skip input layers
-        if isinstance(current_layer, tf.keras.layers.InputLayer):
-            continue
-
-        # If the current layer is either a subclassed layer, functional model or sequential model, we need to extract the
-        # layers from the nested layer and add them to the functional model.
-        if _is_nested_layer(current_layer):
-            new_output_tensor = _handle_nested_layer(current_layer,
-                                                     model_layers_connections,
-                                                     class_names,
-                                                     model_outputs,
-                                                     original_models_last_layer)
-            # If we are at the end of the original model, we want the model_outputs to be the end model outputs
-            if current_layer == original_models_last_layer:
-                _logger.debug("Last layer was a nested layer. "
-                              "Using temp model's output from _handle_nested_layer as model_output")
-                continue
-            model_outputs.clear()
+            input_shape = [shape if shape is not None else 1 for shape in self.prepared_model.input_shape]
+            random_input = np.random.rand(*input_shape).astype(np.float32)
+
+        verbose = logging.DEBUG == _logger.level
+        original_model_output = self.original_model.predict(random_input, verbose=verbose)
+        prepared_model_output = self.prepared_model.predict(random_input, verbose=verbose)
+
+        # Check the outputs of the prepared model and the original model
+        err_msg = """
+        Outputs of prepared model and original model do not match. Since the weights match and params 
+        match, this is likely due to a mismatch in the model's architecture. Specifically, if there is a reuse of a 
+        layer, then the prepared model will not have the same output as the original model. For example, 
+        if a ReLU layer is defined once and then used twice, then the prepared model will only have one ReLU layer 
+        while the original model will have two ReLU layers. Please check the model's architecture to see if there are 
+        any layers that are reused.
+        """
+
+        if isinstance(original_model_output, Dict):
+            original_model_output = [output for output in original_model_output.values()]
+            if len(original_model_output) == 1:
+                original_model_output = original_model_output[0]
+
+        if isinstance(original_model_output, List):
+            for original_output, prepared_output in zip(original_model_output, prepared_model_output):
+                np.testing.assert_array_equal(original_output, prepared_output, err_msg=err_msg)
         else:
-            new_output_tensor = _handle_normal_keras_layer(current_layer, model_layers_connections)
-
-
-        _update_output_tensors_in_model_layers_connections(current_layer,
-                                                           new_output_tensor,
-                                                           model,
-                                                           model_layers_connections,
-                                                           model_outputs)
-    return new_output_tensor
-
-
-def _get_prepared_model(original_model: tf.keras.Model, input_layer: tf.keras.layers.Layer,
-                        class_names: Set[str]) -> tf.keras.Model:
-    """
-    Function to get the prepared model. This function sets up the input layer and calls the helper function to
-    recursively prepare the model.
-
-    :param original_model: The original model to prepare
-    :param input_layer: The input layer to use for the new model
-    :param class_names: The names of the classes that the original model has included for nested layers
-    :return: The prepared model
-    """
-    model_layers_connections = ModelLayerConnections.get_model_layers_connection_properties(original_model)
-    _set_prepared_models_input_layer(original_model, model_layers_connections, input_layer)
+            np.testing.assert_array_equal(original_model_output, prepared_model_output, err_msg=err_msg)
+        _logger.debug("Outputs of prepared model and original model match")
 
-    model_outputs = []
-    _prepare_model_helper(original_model,
-                          class_names,
-                          model_layers_connections,
-                          model_outputs,
-                          original_models_last_layer=original_model.layers[-1])
-
-    # If the model outputs are empty, then we need to get the most recently added output tensor. This is the case
-    # when a model might be sparse and not fully connected or when a Functional model is inside an inherited model.
-    if not model_outputs:
-        _logger.warning("No model outputs found. This usually occurs when a models is made by inheriting from "
-                        "'tf.keras.Model' and placing a Functional model inside. "
-                        "Using most recently added output tensor as prepared models output.")
-        model_outputs = _get_most_recently_added_output_tensor(model_layers_connections)
-
-    return tf.keras.Model(inputs=input_layer, outputs=model_outputs)
-
-
-def _model_has_nested_layers(model: tf.keras.Model) -> bool:
-    """
-    Helper function to check if a model is needed to be prepared or not based on if the model has nested layers such as
-    subclass, functional, or sequential.
-
-    :param model: The model to check
-    :return: If the model needs to be prepared or not
-    """
-    for layer in model.layers:
-        if _is_nested_layer(layer):
-            return True
-    return False
+        _logger.info("Prepared model verified")
 
 
 def prepare_model(original_model: tf.keras.Model,
                   input_layer: Union[tf.keras.layers.InputLayer, List[tf.keras.layers.InputLayer]] = None) \
-                    -> tf.keras.Model:
+        -> tf.keras.Model:
     """
     This function prepares a Keras model before continuing on with AIMET. Specifically, it will convert the model into
     a purely Functional API model and copy over the original models weights.
@@ -633,36 +819,20 @@ def prepare_model(original_model: tf.keras.Model,
     """
 
     # Initial check to see if preparing model is necessary
-    if not _model_has_nested_layers(original_model):
-        _logger.debug("Model does not contain any nested layers. "
-                      "Returning original model after going through 'replace_separable_conv_with_depthwise_pointwise.")
-        model_to_return, _ = replace_separable_conv_with_depthwise_pointwise(original_model)
-        model_to_return, _ = replace_relu6_with_relu(model_to_return)
-        return model_to_return
-
-    _logger.debug("Preparing model for AIMET. Original model architecture")
-    original_model.summary(print_fn=_logger.debug)
-    input_layer = _format_input_layer(original_model, input_layer)
-
-    # Used to fix weight names at end of unwrapping
-    # Originally set to the name of the original model's class in the case that there is an inherited model
-    class_names = _get_class_names_in_model(original_model)
-
-    prepared_model = _get_prepared_model(original_model, input_layer, class_names)
-
-    # Cloning model to remove any references to the original model
-    K.clear_session() # To avoid name conflicts
-    model_to_return = tf.keras.models.clone_model(prepared_model)
-
-    model_to_return.summary(print_fn=_logger.debug)
-
-    # Copying over weights from original model to functional model
-    _logger.debug("Final class_names: %s", class_names)
-    _set_functional_models_weights(original_model, model_to_return, class_names)
-
-    # Extra prepare step to replace Separable Conv's with Depthwise Pointwise pattern if the prepared model
-    # had any in the original models nested layers.
-    model_to_return, _ = replace_separable_conv_with_depthwise_pointwise(model_to_return)
-    model_to_return, _ = replace_relu6_with_relu(model_to_return)
-
-    return model_to_return
+    # pylint: disable=protected-access
+    if not _KerasModelPreparer._model_has_nested_layers(original_model) and \
+            not _KerasModelPreparer._inherits_from_keras_model(original_model):
+        _logger.info("Model does not contain any nested layers. "
+                     "Returning original model after going through "
+                     "'replace_separable_conv_with_depthwise_pointwise' and 'replace_relu6_with_relu.")
+        custom_objects = _KerasModelPreparer._get_models_custom_objects(original_model)
+        prepared_model, _ = replace_relu6_with_relu(original_model, custom_objects=custom_objects)
+        prepared_model, _ = replace_separable_conv_with_depthwise_pointwise(prepared_model,
+                                                                            custom_objects=custom_objects)
+        return prepared_model
+
+    keras_model_preparer = _KerasModelPreparer(original_model, input_layer=input_layer)
+
+    keras_model_preparer.prepare_model()
+
+    return keras_model_preparer.prepared_model
diff --git a/TrainingExtensions/tensorflow/src/python/aimet_tensorflow/keras/quant_sim/qc_quantize_wrapper.py b/TrainingExtensions/tensorflow/src/python/aimet_tensorflow/keras/quant_sim/qc_quantize_wrapper.py
index c874183ebf..9a0c65b491 100755
--- a/TrainingExtensions/tensorflow/src/python/aimet_tensorflow/keras/quant_sim/qc_quantize_wrapper.py
+++ b/TrainingExtensions/tensorflow/src/python/aimet_tensorflow/keras/quant_sim/qc_quantize_wrapper.py
@@ -44,13 +44,13 @@
 from aimet_common.utils import AimetLogger
 from aimet_common.defs import QuantScheme, QuantizationDataType
 import aimet_tensorflow.keras.utils.common as keras_common_utils
-from aimet_tensorflow.keras.model_preparer import _is_keras_or_tensor_input
+from aimet_tensorflow.keras.model_preparer import _KerasModelPreparer
 from aimet_tensorflow.keras.quant_sim.tensor_quantizer import ActivationTensorQuantizer, \
     ParamPerTensorQuantizer, ParamPerChannelQuantizer, StaticGridPerChannelQuantizer
 from aimet_tensorflow.utils.constants import QUANT_ALLOWED_DTYPES
 
 _logger = AimetLogger.get_area_logger(AimetLogger.LogAreas.Quant)
-
+is_tf_or_keras_tensor_input = _KerasModelPreparer._is_tf_or_keras_tensor_input # pylint: disable=protected-access
 
 class QuantizerSettings:
     """ Class holding quantizer settings """
@@ -320,7 +320,7 @@ def call(self, inputs, *args, **kwargs):
             if (self._is_lambda_operator_layer or self._is_a_tf_op_lambda_layer) and len(self.input_quantizers) >= 2:
                 kwargs_keys_for_keras_tensors = [
                     name for name, tensor in kwargs.items()
-                    if _is_keras_or_tensor_input(tensor) or isinstance(tensor, np.ndarray)
+                    if is_tf_or_keras_tensor_input(tensor) or isinstance(tensor, np.ndarray)
                 ]
 
                 # TF functions like tf.concat could have two inputs in List form. But other layers could match
diff --git a/TrainingExtensions/tensorflow/src/python/aimet_tensorflow/keras/utils/model_transform_utils.py b/TrainingExtensions/tensorflow/src/python/aimet_tensorflow/keras/utils/model_transform_utils.py
index 6f236d800d..c176a070b5 100644
--- a/TrainingExtensions/tensorflow/src/python/aimet_tensorflow/keras/utils/model_transform_utils.py
+++ b/TrainingExtensions/tensorflow/src/python/aimet_tensorflow/keras/utils/model_transform_utils.py
@@ -2,7 +2,7 @@
 # =============================================================================
 #  @@-COPYRIGHT-START-@@
 #
-#  Copyright (c) 2022, Qualcomm Innovation Center, Inc. All rights reserved.
+#  Copyright (c) 2022 - 2023, Qualcomm Innovation Center, Inc. All rights reserved.
 #
 #  Redistribution and use in source and binary forms, with or without
 #  modification, are permitted provided that the following conditions are met:
@@ -52,6 +52,19 @@
     )
 
 
+    def _clear_session_before_transformation(transform_func_to_run_after_clearing_session: typing.Callable):
+        """
+        Decorator to clear the backend session before every graph transformation. This is to maintain the naming of the
+        layers and make sure Keras doesn't change them.
+
+        :param transform_func_to_run_after_clearing_session: Transformation function to run after clearing the session.
+        """
+        def wrapper(*args, **kwargs):
+            tf.keras.backend.clear_session()
+            return transform_func_to_run_after_clearing_session(*args, **kwargs)
+        return wrapper
+
+
     def _create_layer_by_config(
             class_name: str, layer_config: typing.Dict
     ) -> tf.keras.layers.Layer:
@@ -75,6 +88,8 @@ class ReplaceRelu6WithRelu(transforms.Transform):
         """
         Transform class for the case like tf.keras.layers.ReLU(max_value=6)
         """
+        def __init__(self, custom_objects: typing.Dict[str, tf.keras.layers.Layer] = None):
+            self.custom_objs = custom_objects
 
         def pattern(self):
             """
@@ -91,11 +106,15 @@ def replacement(self, match_layer):
             replace_layer["name"] = replace_layer["config"]["name"]
             return transforms.LayerNode(replace_layer)
 
+        def custom_objects(self):
+            return self.custom_objs if self.custom_objs else super().custom_objects()
 
     class ReplaceActivationWithRelu(transforms.Transform):
         """
         Transform class for the case like tf.keras.layers.Activation(tf.nn.relu6)
         """
+        def __init__(self, custom_objects: typing.Dict[str, tf.keras.layers.Layer] = None):
+            self.custom_objs = custom_objects
 
         def pattern(self):
             """
@@ -112,14 +131,19 @@ def replacement(self, match_layer):
             replace_layer["name"] = replace_layer["config"]["name"]
             return transforms.LayerNode(replace_layer)
 
+        def custom_objects(self):
+            return self.custom_objs if self.custom_objs else super().custom_objects()
+
 
     class ReplaceFusedRelu6WithFusedRelu(transforms.Transform):
         """
         Transform class for the case like tf.keras.layers.Conv2D(..., activation=tf.nn.relu6)
         Result is still fused such as tf.keras.layers.Conv2D(..., activation=tf.nn.relu)
         """
-        def __init__(self, class_name):
+
+        def __init__(self, class_name, custom_objects: typing.Dict[str, tf.keras.layers.Layer] = None):
             self.class_name = class_name
+            self.custom_objs = custom_objects
 
         def pattern(self):
             """
@@ -141,14 +165,18 @@ def replacement(self, match_layer):
 
             return transforms.LayerNode(replace_layer_config, match_layer.weights)
 
+        def custom_objects(self):
+            return self.custom_objs if self.custom_objs else super().custom_objects()
+
 
     class ReplaceFusedRelu6WithSeparateLayers(transforms.Transform):
         """
         Transform class for the case like tf.keras.layers.Conv2D(..., activation=tf.nn.relu6)
         Result is separated such as layers.Conv2D(..., activation='linear') and layers.ReLU()
         """
-        def __init__(self, class_name):
+        def __init__(self, class_name, custom_objects: typing.Dict[str, tf.keras.layers.Layer] = None):
             self.class_name = class_name
+            self.custom_objs = custom_objects
 
         def pattern(self):
             """
@@ -177,9 +205,13 @@ def replacement(self, match_layer):
                 input_layers=[transforms.LayerNode(replace_layer, match_layer.weights)]
             )
 
+        def custom_objects(self):
+            return self.custom_objs if self.custom_objs else super().custom_objects()
+
 
+    @_clear_session_before_transformation
     def replace_relu6_with_relu(
-            model: tf.keras.Model, remain_fusing: bool = False
+            model: tf.keras.Model, remain_fusing: bool = False, custom_objects: typing.Dict[str, tf.keras.layers.Layer] = None
     ) -> typing.Tuple[tf.keras.Model, typing.Dict]:
         """
         Replace ReLU6 with ReLU in tf.keras.Model
@@ -194,17 +226,17 @@ def replace_relu6_with_relu(
         """
         if remain_fusing:
             transform_list = [
-                ReplaceRelu6WithRelu(),
-                ReplaceActivationWithRelu(),
-                ReplaceFusedRelu6WithFusedRelu("Conv2D"),
-                ReplaceFusedRelu6WithFusedRelu("Dense")
+                ReplaceRelu6WithRelu(custom_objects=custom_objects),
+                ReplaceActivationWithRelu(custom_objects=custom_objects),
+                ReplaceFusedRelu6WithFusedRelu("Conv2D", custom_objects=custom_objects),
+                ReplaceFusedRelu6WithFusedRelu("Dense", custom_objects=custom_objects)
             ]
         else:
             transform_list = [
-                ReplaceRelu6WithRelu(),
-                ReplaceActivationWithRelu(),
-                ReplaceFusedRelu6WithSeparateLayers("Conv2D"),
-                ReplaceFusedRelu6WithSeparateLayers("Dense"),
+                ReplaceRelu6WithRelu(custom_objects=custom_objects),
+                ReplaceActivationWithRelu(custom_objects=custom_objects),
+                ReplaceFusedRelu6WithSeparateLayers("Conv2D", custom_objects=custom_objects),
+                ReplaceFusedRelu6WithSeparateLayers("Dense", custom_objects=custom_objects),
             ]
 
         return model_transformer.ModelTransformer(model, transform_list).transform()
@@ -215,6 +247,9 @@ class ReplaceSeparableConvWithDepthwisePointwise(transforms.Transform):
         Result is separated into two separate layers, layers.DepthwiseConv2D and layers.Conv2D
         """
 
+        def __init__(self, custom_objects: typing.Dict[str, tf.keras.layers.Layer] = None):
+            self.custom_objs = custom_objects
+
         def pattern(self):
             """
             Pattern to transform
@@ -288,11 +323,19 @@ def replacement(self, match_layer):
                 weights=pointwise_layer_weights,
                 input_layers=[transforms.LayerNode(depthwise_layer_config, weights=depthwise_layer_weights)])
 
-    def replace_separable_conv_with_depthwise_pointwise(model: tf.keras.Model) -> typing.Tuple[tf.keras.Model, typing.Dict]:
+        def custom_objects(self):
+            return self.custom_objs if self.custom_objs else super().custom_objects()
+
+
+    @_clear_session_before_transformation
+    def replace_separable_conv_with_depthwise_pointwise(model: tf.keras.Model,
+                                                        custom_objects: typing.Dict[str, tf.keras.layers.Layer] = None) -> \
+                                                            typing.Tuple[tf.keras.Model, typing.Dict]:
         """
         Replace SeparableConv2D with DepthwiseConv2D and Conv2D in tf.keras.Model
 
         :param model: tf.keras.Model
+        :param custom_objects: Models custom_objects.
         """
-        transform_list = [ReplaceSeparableConvWithDepthwisePointwise()]
+        transform_list = [ReplaceSeparableConvWithDepthwisePointwise(custom_objects=custom_objects)]
         return model_transformer.ModelTransformer(model, transform_list).transform()


Test patch:
diff --git a/TrainingExtensions/tensorflow/test/python/eager/test_keras_utils.py b/TrainingExtensions/tensorflow/test/python/eager/test_keras_utils.py
index 4a5f50d761..7d953a38b0 100644
--- a/TrainingExtensions/tensorflow/test/python/eager/test_keras_utils.py
+++ b/TrainingExtensions/tensorflow/test/python/eager/test_keras_utils.py
@@ -66,43 +66,6 @@ def conv_functional():
     return model
 
 
-# Not used for testing at the moment. This is placed here for future testing.
-class ConvTimesThree(tf.keras.layers.Layer):
-    def __init__(self, **kwargs):
-        super(ConvTimesThree, self).__init__(**kwargs)
-        self.conv = tf.keras.layers.Conv2D(32,
-                                           kernel_size=(3, 3),
-                                           activation='relu',
-                                           name='class_conv')
-        self.conv_transpose = tf.keras.layers.Conv2DTranspose(64,
-                                                              kernel_size=(3, 3),
-                                                              activation='relu',
-                                                              name='class_conv_transpose')
-        self.depth_conv = tf.keras.layers.DepthwiseConv2D(depth_multiplier=1,
-                                                          kernel_size=(3, 3),
-                                                          activation='relu',
-                                                          name='class_conv_depth')
-
-    def call(self, x):
-        x = self.conv(x)
-        x = self.conv_transpose(x)
-        x = self.depth_conv(x)
-        return x
-
-
-# See comment above ConvTimesThree Class
-def conv_sub_class():
-    input_shape = (128, 28, 28, 1)
-    inp = tf.keras.Input(shape=input_shape[1:])
-    x = ConvTimesThree()(inp)
-    x = tf.keras.layers.Flatten()(x)
-    x = tf.keras.layers.Dropout(0.5)(x, training=False)
-    x = tf.keras.layers.Dense(10, activation="softmax")(x)
-
-    model = tf.keras.Model(inputs=inp, outputs=x, name='conv_classes')
-    return model
-
-
 def test_replace_middle_layers():
     with tempfile.TemporaryDirectory() as tmp_dir:
         saved_model_dir = os.path.join(tmp_dir, 'saved_model')
@@ -359,13 +322,6 @@ def test_convert_h5_to_pb_functional_model():
     check_conversion_tensor_names(conv_functional())
 
 
-@pytest.mark.skip(
-    reason="Subclassed Keras models are not currently supported. "
-           "Created for future testing.")
-def test_convert_h5_to_pb_subclass_model():
-    check_conversion_tensor_names(conv_sub_class())
-
-
 @pytest.mark.skip(reason="This test takes a long time. Only used during development.")
 def test_convert_h5_to_pb_pretrained_keras():
     model = tf.keras.applications.ResNet50(weights="imagenet",
diff --git a/TrainingExtensions/tensorflow/test/python/eager/test_model_preparer.py b/TrainingExtensions/tensorflow/test/python/eager/test_model_preparer.py
index 3422d5ed07..9158f0055a 100644
--- a/TrainingExtensions/tensorflow/test/python/eager/test_model_preparer.py
+++ b/TrainingExtensions/tensorflow/test/python/eager/test_model_preparer.py
@@ -38,46 +38,40 @@
 """ Unit tests for keras model preparer """
 import logging
 import os
+import tempfile
+from typing import List
+
 import pytest
 import numpy as np
 import tensorflow as tf
+
 from aimet_common.utils import AimetLogger
+from aimet_tensorflow.keras.quantsim import QuantizationSimModel
 from aimet_tensorflow.keras.connectedgraph import ConnectedGraph
-from aimet_tensorflow.keras.model_preparer import prepare_model, _get_original_models_weights_in_functional_model_order
-from aimet_common.connected_graph.connectedgraph_utils import get_all_input_ops, get_all_output_ops
-
-os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
-AimetLogger.set_level_for_all_areas(logging.DEBUG)
+from aimet_tensorflow.keras.model_preparer import prepare_model, _KerasModelPreparer
 
+get_models_custom_objects = _KerasModelPreparer._get_models_custom_objects
+from aimet_common.connected_graph.connectedgraph_utils import get_all_input_ops, get_all_output_ops
 
-def conv_functional():
-    input_shape = (128, 28, 28, 1)
-    inp = tf.keras.Input(shape=input_shape[1:])
-    x = tf.keras.layers.Conv2D(32, kernel_size=(3, 3), activation="relu")(inp)
-    x = tf.keras.layers.Conv2DTranspose(
-        32, kernel_size=(3, 3), activation="relu")(x)
-    x = tf.keras.layers.DepthwiseConv2D(
-        depth_multiplier=1, kernel_size=(3, 3), activation='relu')(x)
-    x = tf.keras.layers.Flatten()(x)
-    x = tf.keras.layers.Dropout(0.5, trainable=False)(x)
-    x = tf.keras.layers.Dense(10, activation="softmax")(x)
+from test_models_keras import resnet_34, tiny_conv_net
 
-    model = tf.keras.Model(inputs=inp, outputs=x, name='conv_functional')
-    return model
+os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
+AimetLogger.set_level_for_all_areas(logging.DEBUG)
 
 
+# Begin of Subclass Models to for Testing
 class TwoConvs(tf.keras.layers.Layer):
     def __init__(self, **kwargs):
         super(TwoConvs, self).__init__(**kwargs)
         self.conv = tf.keras.layers.Conv2D(32,
                                            kernel_size=(3, 3),
-                                           activation='relu',
-                                           name='class_conv')
+                                           activation="relu",
+                                           name="class_conv")
 
         self.conv_transpose = tf.keras.layers.Conv2DTranspose(64,
                                                               kernel_size=(3, 3),
-                                                              activation='relu',
-                                                              name='class_conv_transpose')
+                                                              activation="relu",
+                                                              name="class_conv_transpose")
 
     def call(self, x, **kwargs):
         x = self.conv(x)
@@ -87,29 +81,17 @@ def call(self, x, **kwargs):
 
 class ConvTimesThree(tf.keras.layers.Layer):
     def __init__(self, **kwargs):
-
         super(ConvTimesThree, self).__init__(**kwargs)
         self.depth_conv = tf.keras.layers.DepthwiseConv2D(depth_multiplier=1,
                                                           kernel_size=(3, 3),
-                                                          activation='relu',
-                                                          name='class_conv_depth')
+                                                          activation="relu",
+                                                          name="class_conv_depth")
         self.two_convs = TwoConvs()
 
     def call(self, x, **kwargs):
         return self.depth_conv(self.two_convs(x))
 
 
-def conv_sub_class():
-    input_shape = (128, 28, 28, 1)
-    inp = tf.keras.Input(batch_shape=input_shape)
-    x = ConvTimesThree()(inp)
-    x = tf.keras.layers.Flatten()(x)
-    x = tf.keras.layers.Dropout(0.5)(x, training=False)
-    x = tf.keras.layers.Dense(10, activation="softmax")(x)
-
-    model = tf.keras.Model(inputs=inp, outputs=x, name='conv_classes')
-    return model
-
 ###########################################################################################################
 #                                                                                                         #
 # Below models are based on Deep Learning with Python by Francois Chollet Second Edition (page 182 - 185) #
@@ -118,7 +100,6 @@ def conv_sub_class():
 
 # Only Subclassing
 class CustomerTicketModel(tf.keras.Model):
-
     def __init__(self, num_departments):
         super().__init__()
         self.concat_layer = tf.keras.layers.Concatenate()
@@ -137,9 +118,9 @@ def call(self, inputs, **kwargs):
         department = self.department_classifier(features)
         return priority, department
 
+
 # Functional model that includes subclassed layers
 class Classifier(tf.keras.Model):
-
     def __init__(self, num_classes=4):
         super().__init__()
         if num_classes == 2:
@@ -154,53 +135,50 @@ def call(self, inputs, **kwargs):
         return self.dense(inputs)
 
 
-def functional_model_with_subclassed_layers():
-    inputs = tf.keras.layers.Input(shape=(3,))
-    features = tf.keras.layers.Dense(64, activation="relu")(inputs)
-    outputs = Classifier(num_classes=10)(features)
-    model = tf.keras.Model(inputs=inputs, outputs=outputs)
-    return model
+# Linear class source: https://keras.io/guides/making_new_layers_and_models_via_subclassing/
+class Linear(tf.keras.layers.Layer):
+    def __init__(self, units=32, **kwargs):
+        super().__init__(**kwargs)
+        self.units = units
+
+    def build(self, input_shape):
+        self.w = self.add_weight(
+            name="weight",
+            shape=(input_shape[-1], self.units),
+            initializer="random_normal",
+            trainable=True,
+        )
 
+        self.b = self.add_weight(
+            name="bias",
+            shape=(self.units,),
+            initializer="random_normal",
+            trainable=True
+        )
 
-# Subclass model that includes functional layers
-def subclass_model_with_functional_layers():
-    inputs = tf.keras.Input(shape=(64,))
-    outputs = tf.keras.layers.Dense(1, activation="sigmoid")(inputs)
-    binary_classifier = tf.keras.Model(inputs=inputs, outputs=outputs)
+    def call(self, inputs):
+        return tf.matmul(inputs, self.w) + self.b
 
-    class MyFunctionalModel(tf.keras.Model):
-        def __init__(self):
-            super().__init__(name='my_functional_model')
-            self.dense = tf.keras.layers.Dense(64, activation="relu")
-            self.classifier = binary_classifier
+    def get_config(self):
+        config = super().get_config()
+        config.update({"units": self.units})
+        return config
 
-        def call(self, inputs, **kwargs):
-            features = self.dense(inputs)
-            return self.classifier(features)
 
-    model = MyFunctionalModel()
-    return model
+class KerasDefinedAndAcceptableDefined(tf.keras.layers.Layer):
 
-# Layer with multiple math operations in call
-class MultiMathOperations(tf.keras.layers.Layer):
-    def __init__(self, **kwargs):
-        super(MultiMathOperations, self).__init__(**kwargs)
-        self.r = tf.constant(np.random.rand(1, 1), dtype=tf.float32)
-        self.g = tf.constant(np.random.rand(1, 1), dtype=tf.float32)
-        self.b = tf.constant(np.random.rand(1, 1), dtype=tf.float32)
+    def __init__(self):
+        super().__init__(name="keras_defined_and_acceptable_defined")
+        self.conv = tf.keras.layers.Conv2D(32, kernel_size=(3, 3), )
+        self.linear = Linear()
+        self.relu = tf.keras.layers.ReLU()
 
     def call(self, inputs, **kwargs):
-        r = tf.math.multiply(inputs[:, :, :, 0:1], self.r)
-        g = tf.math.multiply(inputs[:, :, :, 1:2], self.g)
-        b = tf.math.multiply(inputs[:, :, :, 2:3], self.b)
-        return tf.concat([r, g, b], axis=3)
-
+        x = self.conv(inputs)
+        x = self.linear(x)
+        x = self.relu(x)
+        return x
 
-    #####################################################
-    #                                                   #
-    #    QA models for testing the conversion script    #
-    #                                                   #
-    #####################################################
 
 # Text classification Transformer model from https://keras.io/examples/nlp/text_classification_with_transformer/
 class TransformerBlock(tf.keras.layers.Layer):
@@ -208,7 +186,7 @@ def __init__(self, embed_dim, num_heads, ff_dim, rate=0.1):
         super(TransformerBlock, self).__init__()
         self.att = tf.keras.layers.MultiHeadAttention(num_heads=num_heads, key_dim=embed_dim)
         self.ffn = tf.keras.Sequential(
-            [tf.keras.layers.Dense(ff_dim, activation="relu"), tf.keras.layers.Dense(embed_dim),]
+            [tf.keras.layers.Dense(ff_dim, activation="relu"), tf.keras.layers.Dense(embed_dim), ]
         )
         self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
         self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
@@ -224,7 +202,6 @@ def call(self, inputs, training, **kwargs):
         return self.layernorm2(out1 + ffn_output)
 
 
-
 class TokenAndPositionEmbedding(tf.keras.layers.Layer):
     def __init__(self, maxlen, vocab_size, embed_dim):
         super(TokenAndPositionEmbedding, self).__init__()
@@ -240,298 +217,359 @@ def call(self, x, **kwargs):
         return x
 
 
+# Layer with multiple math operations in call
+class MultiMathOperations(tf.keras.layers.Layer):
+    def __init__(self, **kwargs):
+        super(MultiMathOperations, self).__init__(**kwargs)
+        self.r = tf.constant(np.random.rand(1, 1), dtype=tf.float32)
+        self.g = tf.constant(np.random.rand(1, 1), dtype=tf.float32)
+        self.b = tf.constant(np.random.rand(1, 1), dtype=tf.float32)
+
+    def call(self, inputs, **kwargs):
+        r = tf.math.multiply(inputs[:, :, :, 0:1], self.r)
+        g = tf.math.multiply(inputs[:, :, :, 1:2], self.g)
+        b = tf.math.multiply(inputs[:, :, :, 2:3], self.b)
+        return tf.concat([r, g, b], axis=3)
+
+
+# Layer with multiple inputs
+class TestMultiInput(tf.keras.layers.Layer):
+    def __init__(self, **kwargs) -> None:
+        super().__init__(name="test_multi_input")
+
+    def call(self, inputs, **kwargs):
+        input1 = inputs[0]
+        input2 = inputs[1]
+        return input1 + input2
+
+
+# Layer with multiple outputs
+class TestMultiOut(tf.keras.layers.Layer):
+    def __init__(self, **kwargs) -> None:
+        super().__init__(name="test_multi_out")
+
+    def call(self, inputs, **kwargs):
+        out1 = inputs * 2.0
+        out2 = inputs * 3.0
+        return [out1, out2]
+
+
+class TestModelPreparer:
+    """ Class for Testing aimet_tensorflow.keras Model Preparer"""
+
+    # Model Getters
+    @staticmethod
+    def get_conv_sub_class():
+        input_shape = (128, 28, 28, 1)
+        inp = tf.keras.Input(batch_shape=input_shape)
+        x = ConvTimesThree()(inp)
+        x = tf.keras.layers.Flatten()(x)
+        x = tf.keras.layers.Dropout(0.5)(x, training=False)
+        x = tf.keras.layers.Dense(10, activation="softmax")(x)
+
+        model = tf.keras.Model(inputs=inp, outputs=x, name="conv_classes")
+        return model
+
+    @staticmethod
+    def get_functional_model_with_subclassed_layers():
+        inputs = tf.keras.layers.Input(shape=(3,))
+        features = tf.keras.layers.Dense(64, activation="relu")(inputs)
+        outputs = Classifier(num_classes=10)(features)
+        model = tf.keras.Model(inputs=inputs, outputs=outputs)
+        return model
+
+    @staticmethod
+    def get_subclass_model_with_functional_layers():
+        inputs = tf.keras.Input(shape=(64,))
+        outputs = tf.keras.layers.Dense(1, activation="sigmoid")(inputs)
+        binary_classifier = tf.keras.Model(inputs=inputs, outputs=outputs)
+
+        class MyFunctionalModel(tf.keras.Model):
+            def __init__(self):
+                super().__init__(name="my_functional_model")
+                self.dense = tf.keras.layers.Dense(64, activation="relu")
+                self.classifier = binary_classifier
+
+            def call(self, inputs, **kwargs):
+                features = self.dense(inputs)
+                return self.classifier(features)
+
+        model = MyFunctionalModel()
+        return model
+
+    @staticmethod
+    def get_keras_text_classification_example_model_and_data_input():
+        # Download and prepare dataset
+        vocab_size = 20000  # Only consider the top 20k words
+        maxlen = 200  # Only consider the first 200 words of each movie review
+
+        random_input = np.random.random((10, 200))
+
+        # Create classifier model using transformer layer
+        embed_dim = 32  # Embedding size for each token
+        num_heads = 2  # Number of attention heads
+        ff_dim = 32  # Hidden layer size in feed forward network inside transformer
+
+        inputs = tf.keras.layers.Input(shape=(maxlen,))
+        embedding_layer = TokenAndPositionEmbedding(maxlen, vocab_size, embed_dim)
+        x = embedding_layer(inputs)
+        transformer_block = TransformerBlock(embed_dim, num_heads, ff_dim)
+        x = transformer_block(x)
+        x = tf.keras.layers.GlobalAveragePooling1D()(x)
+        x = tf.keras.layers.Dropout(0.1)(x)
+        x = tf.keras.layers.Dense(20, activation="relu")(x)
+        x = tf.keras.layers.Dropout(0.1)(x)
+        outputs = tf.keras.layers.Dense(2, activation="softmax")(x)
+
+        return tf.keras.Model(inputs=inputs, outputs=outputs, name="text_classification")
+
+    @staticmethod
+    def get_multi_math_op_model():
+        input_layer = tf.keras.Input(shape=(32, 32, 1))
+        output = MultiMathOperations()(input_layer)
+        return tf.keras.Model(inputs=input_layer, outputs=output)
+
+    @staticmethod
+    def get_model_with_multiple_inputs():
+        input1 = tf.keras.Input(shape=(28, 28, 1), name="input_1")
+        input2 = tf.keras.Input(shape=(28, 28, 1), name="input_2")
+        x = tf.keras.layers.Conv2D(16, 3, activation="relu")(input1)
+        y = tf.keras.layers.Conv2D(16, 3, activation="relu")(input2)
+        outputs = TestMultiInput()([x, y])
+
+        return tf.keras.Model(inputs=[input1, input2], outputs=outputs, name="multi_input")
+
+    @staticmethod
+    def get_model_with_multiple_outputs(use_lambdas):
+        inputs = tf.keras.Input(shape=(28, 28, 1), name="img")
+        x = tf.keras.layers.Conv2D(16, 3, activation="relu")(inputs)
+        if use_lambdas:
+            x0, x1 = tf.split(x, [8, 8], 3)
+            x1a = x1 * 2.0
+            outputs = tf.concat([x0, x1a], 3)
+        else:
+            outputs = TestMultiOut()(x)
+
+        return tf.keras.Model(inputs=inputs, outputs=outputs, name=f"multi_output_with_lambda_{use_lambdas}")
+
+    @staticmethod
+    def get_model_with_acceptable_subclass_layers():
+        inputs = tf.keras.Input(shape=(32,))
+        x = Linear()(inputs)
+        x = tf.keras.layers.ReLU()(x)
+        x = Linear()(x)
+        outputs = tf.keras.layers.ReLU()(x)
+
+        return tf.keras.Model(inputs=inputs, outputs=outputs, name="acceptable_subclasses")
+
+    @staticmethod
+    def get_model_with_keras_and_acceptable_subclass_layers_defined():
+        inputs = tf.keras.Input(shape=(28, 28, 1,))
+        outputs = KerasDefinedAndAcceptableDefined()(inputs)
+
+        return tf.keras.Model(inputs=inputs, outputs=outputs)
+
+    # Helper Functions
+    def get_model_and_set_random_input(self, get_model_func):
+        original_model = get_model_func()
+        if isinstance(original_model.input_shape, List):
+            self.random_input = [
+                np.random.rand(1, *input_shape[1:]) for input_shape in original_model.input_shape
+            ]
+        else:
+            self.random_input = np.random.rand(1, *original_model.input_shape[1:])
+        return original_model
+
+    @staticmethod
+    def get_model_inherited_from_Model_without_config() -> tf.keras.Model:
+        class LayerWithoutConfig(tf.keras.Model):  # Note: Inherits from Model
+            def __init__(self, value_to_save_in_config, *args, **kwargs):
+                super().__init__(*args, **kwargs)
+                self.dense = tf.keras.layers.Dense(value_to_save_in_config, 'relu')
+
+            def call(self, inputs, training=None, mask=None):
+                return self.dense(inputs)
+
+
+        class MyTempModel(tf.keras.Model):
+            def __init__(self):
+                super(MyTempModel, self).__init__()
+                self.layer_without_config = LayerWithoutConfig(3)
+
+            def call(self, inputs, training=None, mask=None):
+                return self.layer_without_config(inputs)
+
+        model = MyTempModel()
+        model.build(input_shape=(1, 28, 28, 3))
+        return model
+
+    # Tests
+    @pytest.fixture(scope="function", autouse=True)
+    def pass_prepared_model_through_quantsim(self):
+        tf.keras.backend.clear_session()
+        yield
+        if not hasattr(self, "prepared_model"):
+            return
+
+        # TODO: multi_output_with_lambda currently fails in ConnectedGraph -> AIMET-2682
+        if self.prepared_model.name.startswith("multi_output_with_lambda"):
+            return
+
+        with tempfile.TemporaryDirectory() as temp_dir:
+            sim = QuantizationSimModel(self.prepared_model)
+            sim.compute_encodings(lambda model, _: model(self.random_input), None)
+            sim.export(temp_dir, self.prepared_model.name, custom_objects=self.__dict__.get("custom_objects"))
+        del self.prepared_model
+
+    @pytest.mark.skip("Need to verify how functional model should look")
+    def test_full_subclass_to_functional(self):
+        vocabulary_size = 10000
+        num_tags = 100
+        num_departments = 4
+        num_samples = 1280
+
+        title_data = np.random.randint(0, 2, size=(1, num_samples, vocabulary_size))
+        text_body_data = np.random.randint(0, 2, size=(1, num_samples, vocabulary_size))
+        tags_data = np.random.randint(0, 2, size=(1, num_samples, num_tags))
+
+        model = CustomerTicketModel(num_departments=num_departments)
+        _ = model({"title": title_data,
+                   "text_body": text_body_data,
+                   "tags": tags_data})
+        # Since this model is fully subclassed, specifically at the beginning, we call prepare model with
+        # the inputs to have Keras symoblic tensor fit the rest of the layers correctly.
+        input_layers = [tf.keras.Input(shape=(num_samples, vocabulary_size,), name="title"),
+                        tf.keras.Input(shape=(num_samples, vocabulary_size,), name="text_body"),
+                        tf.keras.Input(shape=(num_samples, num_tags,), name="tags")]
+        _ = prepare_model(model, input_layers)
+
+    @pytest.mark.parametrize("model_func_str", [
+        "get_multi_math_op_model",
+        "get_model_with_multiple_inputs",
+        "get_keras_text_classification_example_model_and_data_input",
+        "get_model_with_keras_and_acceptable_subclass_layers_defined",
+        "get_model_with_multiple_outputs"
+    ])
+    def test_models_with_common_flow(self, model_func_str):
+        model_func = self.__getattribute__(model_func_str)
+        if model_func_str == "get_model_with_multiple_outputs":
+            model_func = lambda: self.__getattribute__(model_func_str)(False)
+        original_model = self.get_model_and_set_random_input(model_func)
+
+        # NOTE: Verification of the model occurs in the prepare_model function.
+        self.prepared_model = prepare_model(original_model)
+        self.custom_objects = get_models_custom_objects(self.prepared_model)
+
+    def test_functional_model_with_subclassed_layers_to_functional(self):
+        original_model = self.get_model_and_set_random_input(self.get_functional_model_with_subclassed_layers)
+
+        self.prepared_model = prepare_model(original_model)
+
+        connected_graph = ConnectedGraph(self.prepared_model)
+        assert "Unknown" not in connected_graph.ordered_ops
+
+        product_dict = connected_graph.get_all_products()
+        assert "Gemm_0_to_Gemm_1" in product_dict
+
+        input_ops = get_all_input_ops(connected_graph)
+        assert len(input_ops) == 1
+        assert input_ops[0].get_module() == self.prepared_model.layers[1]
+        assert isinstance(input_ops[0].get_module(), type(original_model.layers[1]))
+
+        output_ops = get_all_output_ops(connected_graph)
+        assert len(output_ops) == 1
+        assert output_ops[0].get_module() == self.prepared_model.layers[-1]
+        assert isinstance(output_ops[0].get_module(), type(original_model.layers[-1].dense))
+
+    def test_input_layer_missing(self):
+        model = self.get_subclass_model_with_functional_layers()
+        input_shape = (32, 64)
+        self.random_input = np.random.rand(*input_shape)
+        _ = model(self.random_input)
+        with pytest.raises(ValueError):
+            prepare_model(model)
+
+    def test_subclass_model_with_subclassed_layers_to_functional(self):
+        original_model = self.get_subclass_model_with_functional_layers()
+        input_shape = (32, 64)
+        self.random_input = np.random.rand(*input_shape)
+        _ = original_model(self.random_input)
+
+        self.prepared_model = prepare_model(original_model, tf.keras.Input(shape=input_shape[1:]))
+
+        connected_graph = ConnectedGraph(self.prepared_model)
+        assert "Unknown" not in connected_graph.ordered_ops
+
+        product_dict = connected_graph.get_all_products()
+        assert "Gemm_0_to_Gemm_1" in product_dict
+
+        input_ops = get_all_input_ops(connected_graph)
+        assert len(input_ops) == 1
+        assert input_ops[0].get_module() == self.prepared_model.layers[1]
+        assert isinstance(input_ops[0].get_module(), type(original_model.layers[0]))
+
+        output_ops = get_all_output_ops(connected_graph)
+        assert len(output_ops) == 1
+        assert output_ops[0].get_module() == self.prepared_model.layers[-1]
+        assert isinstance(output_ops[0].get_module(), type(original_model.layers[-1].layers[-1]))
+
+    def test_conv_times_three_subclass_to_functional(self):
+        original_model = self.get_model_and_set_random_input(self.get_conv_sub_class)
+
+        self.prepared_model = prepare_model(original_model)
+
+        connected_graph = ConnectedGraph(self.prepared_model)
+        assert "Unknown" not in connected_graph.ordered_ops
+
+        product_dict = connected_graph.get_all_products()
+        assert "Conv_0_to_ConvTranspose_1" in product_dict
+        assert "ConvTranspose_1_to_Conv_2" in product_dict
+
+        input_ops = get_all_input_ops(connected_graph)
+        assert len(input_ops) == 1
+        assert input_ops[0].get_module() == self.prepared_model.layers[1]
+        assert isinstance(input_ops[0].get_module(), type(original_model.layers[1].two_convs.conv))
+
+        output_ops = get_all_output_ops(connected_graph)
+        assert len(output_ops) == 1
+        assert output_ops[0].get_module() == self.prepared_model.layers[-1]
+        assert isinstance(output_ops[0].get_module(), type(original_model.layers[-1]))
+
+    def test_non_nested_layered_model(self):
+        original_model = self.get_model_and_set_random_input(tiny_conv_net)
+
+        self.prepared_model = prepare_model(original_model)
+
+        assert original_model.get_config() == self.prepared_model.get_config(), \
+            "Prepare model did not give back the original model. This model does not need to be prepared."
+
+    def test_multi_output_with_lambdas(self):
+        original_model = self.get_model_and_set_random_input(lambda: self.get_model_with_multiple_outputs(True))
+
+        # NOTE: Verification of the model occurs in the prepare_model function.
+        self.prepared_model = prepare_model(original_model)
+        assert original_model.get_config() == self.prepared_model.get_config(), \
+            "The original model does not contain any nested layers. The original model should be returned."
+
+    def test_model_with_acceptable_subclass_layers(self):
+        original_model = self.get_model_and_set_random_input(self.get_model_with_acceptable_subclass_layers)
+
+        self.prepared_model = prepare_model(original_model)
+        assert original_model.get_config() == self.prepared_model.get_config(), \
+            "The original model contains acceptable subclass layers. The original model should be returned."
+        self.custom_objects = get_models_custom_objects(self.prepared_model)
+
+    def test_model_inherits_from_Model_without_get_config_defined(self):
+        original_model = self.get_model_inherited_from_Model_without_config()
+
+        with pytest.raises(TypeError):
+            prepare_model(original_model, input_layer=tf.keras.Input(shape=[28, 28, 3]))
+
+    def test_resnet_34(self):
+        original_model = resnet_34()
+        input_shape = (480, 480, 3)
+        self.random_input = np.random.rand(1, *input_shape)
+        _ = original_model(self.random_input)
 
-def get_keras_text_classification_example_model_and_data_input():
-    # Download and prepare dataset
-    vocab_size = 20000  # Only consider the top 20k words
-    maxlen = 200  # Only consider the first 200 words of each movie review
-
-    random_input = np.random.random((10, 200))
-
-    # Create classifier model using transformer layer
-    embed_dim = 32  # Embedding size for each token
-    num_heads = 2  # Number of attention heads
-    ff_dim = 32  # Hidden layer size in feed forward network inside transformer
-
-    inputs = tf.keras.layers.Input(shape=(maxlen,))
-    embedding_layer = TokenAndPositionEmbedding(maxlen, vocab_size, embed_dim)
-    x = embedding_layer(inputs)
-    transformer_block = TransformerBlock(embed_dim, num_heads, ff_dim)
-    x = transformer_block(x)
-    x = tf.keras.layers.GlobalAveragePooling1D()(x)
-    x = tf.keras.layers.Dropout(0.1)(x)
-    x = tf.keras.layers.Dense(20, activation="relu")(x)
-    x = tf.keras.layers.Dropout(0.1)(x)
-    outputs = tf.keras.layers.Dense(2, activation="softmax")(x)
-
-    model = tf.keras.Model(inputs=inputs, outputs=outputs)
-    _ = model(random_input)
-    # Train and Evaluate
-    return model, random_input
-
-
-    ################################
-    #                              #
-    #  Start of testing functions  #
-    #                              #
-    ################################
-
-def compare_weights(original_weights, functional_weights):
-    """
-    Helper function to compare the weights of two models. This function is used to test the conversion script.
-    :param original_weights: the original model's weights 
-    :param functional_weights: the model's weights that was converted from the original model
-    """
-
-    for i, _ in enumerate(original_weights):
-        np.testing.assert_array_equal(original_weights[i], functional_weights[i])
-
-
-def verify_functional_model(functional_model, original_model,
-                            random_input, number_of_layers_in_model):
-    """
-    Helper function to verify that the functional model is ready for AIMET. 
-    This function is used to test the conversion script.
-    :param functional_model: the functional model that was converted from the original model
-    :param original_model: the original model
-    :param random_input: a random input to the model
-    """
-    # Verify the functional model produces the same output as the original model
-    try:
-        np.testing.assert_array_equal(functional_model(random_input).numpy(), original_model(random_input).numpy())
-    except AttributeError:
-        tf.assert_equal(functional_model(random_input), original_model(random_input))
-    assert len(functional_model.layers) == number_of_layers_in_model, \
-        f"Expected {number_of_layers_in_model} layers in functional model, but got {len(functional_model.layers)}"
-
-
-@pytest.mark.skip("Need to verify how functional model should look")
-def test_full_subclass_to_functional():
-    vocabulary_size = 10000
-    num_tags = 100
-    num_departments = 4
-    num_samples = 1280
-
-    title_data = np.random.randint(0, 2, size=(1, num_samples, vocabulary_size))
-    text_body_data = np.random.randint(0, 2, size=(1, num_samples, vocabulary_size))
-    tags_data = np.random.randint(0, 2, size=(1, num_samples, num_tags))
-
-    model = CustomerTicketModel(num_departments=num_departments)
-    _ = model({"title": title_data,
-               "text_body": text_body_data,
-               "tags": tags_data})
-    # Since this model is fully subclassed, specifically at the beginning, we call prepare model with
-    # the inputs to have Keras symoblic tensor fit the rest of the layers correctly.
-    input_layers = [tf.keras.Input(shape=(num_samples, vocabulary_size,), name="title"),
-                    tf.keras.Input(shape=(num_samples, vocabulary_size,), name="text_body"),
-                    tf.keras.Input(shape=(num_samples, num_tags,), name="tags")]
-    functional_model = prepare_model(model, input_layers)
-    assert functional_model.count_params() == model.count_params()
-
-
-def test_functional_model_with_subclassed_layers_to_functional():
-    original_model = functional_model_with_subclassed_layers()
-    random_input = np.random.rand(32, 3)
-    _ = original_model(random_input)
-
-    functional_model = prepare_model(original_model)
-    assert functional_model.count_params() == original_model.count_params()
-
-    compare_weights(original_model.get_weights(), functional_model.get_weights())
-    np.testing.assert_array_equal(functional_model(random_input).numpy(), original_model(random_input).numpy())
-    verify_functional_model(functional_model,
-                            original_model,
-                            random_input,
-                            number_of_layers_in_model=3)
-
-    connected_graph = ConnectedGraph(functional_model)
-    assert "Unknown" not in connected_graph.ordered_ops
-
-    product_dict = connected_graph.get_all_products()
-    assert "Gemm_0_to_Gemm_1" in product_dict
-
-    input_ops = get_all_input_ops(connected_graph)
-    assert len(input_ops) == 1
-    assert input_ops[0].get_module() == functional_model.layers[1]
-    assert isinstance(input_ops[0].get_module(), type(original_model.layers[1]))
-
-    output_ops = get_all_output_ops(connected_graph)
-    assert len(output_ops) == 1
-    assert output_ops[0].get_module() == functional_model.layers[-1]
-    assert isinstance(output_ops[0].get_module(), type(original_model.layers[-1].dense))
-
-
-def test_input_layer_missing():
-    model = subclass_model_with_functional_layers()
-    input_shape = (32, 64)
-    random_input = np.random.rand(*input_shape)
-    _ = model(random_input)
-    with pytest.raises(ValueError):
-        prepare_model(model)
-
-
-def test_subclass_model_with_subclassed_layers_to_functional():
-    original_model = subclass_model_with_functional_layers()
-    input_shape = (32, 64)
-    random_input = np.random.rand(*input_shape)
-    _ = original_model(random_input)
-
-    functional_model = prepare_model(original_model, tf.keras.Input(shape=input_shape[1:]))
-    assert functional_model.count_params() == original_model.count_params()
-    compare_weights(original_model.get_weights(), functional_model.get_weights())
-    verify_functional_model(functional_model,
-                            original_model,
-                            random_input,
-                            number_of_layers_in_model=3)
-
-    connected_graph = ConnectedGraph(functional_model)
-    assert "Unknown" not in connected_graph.ordered_ops
-
-    product_dict = connected_graph.get_all_products()
-    assert "Gemm_0_to_Gemm_1" in product_dict
-
-    input_ops = get_all_input_ops(connected_graph)
-    assert len(input_ops) == 1
-    assert input_ops[0].get_module() == functional_model.layers[1]
-    assert isinstance(input_ops[0].get_module(), type(original_model.layers[0]))
-
-    output_ops = get_all_output_ops(connected_graph)
-    assert len(output_ops) == 1
-    assert output_ops[0].get_module() == functional_model.layers[-1]
-    assert isinstance(output_ops[0].get_module(), type(original_model.layers[-1].layers[-1]))
-
-
-def test_conv_times_three_subclass_to_functional():
-    original_model = conv_sub_class()
-    input_shape = (32, 28, 28, 1)
-    random_input = np.random.rand(*input_shape)
-    _ = original_model(random_input)
-
-    functional_model = prepare_model(original_model)
-    assert functional_model.count_params() == original_model.count_params()
-
-    # NOTE: Since ConvTimesThree has the internal layers out of order compared to the call method,
-    # the weights are not in the order of what the actual architecture is (this is a Keras design).
-    # Therefore, we get the original model's weights and sort them in the order of the actual
-    # architecture and use those weights to compare to the functional model's weights.
-    model_weights_in_correct_order = _get_original_models_weights_in_functional_model_order(
-        original_model, functional_model, class_names={"conv_times_three", "two_convs"})
-
-    compare_weights(model_weights_in_correct_order, functional_model.get_weights())
-    verify_functional_model(functional_model,
-                            original_model,
-                            random_input,
-                            number_of_layers_in_model=7)
-
-    connected_graph = ConnectedGraph(functional_model)
-    assert "Unknown" not in connected_graph.ordered_ops
-
-    product_dict = connected_graph.get_all_products()
-    assert "Conv_0_to_ConvTranspose_1" in product_dict
-    assert "ConvTranspose_1_to_Conv_2" in product_dict
-
-    input_ops = get_all_input_ops(connected_graph)
-    assert len(input_ops) == 1
-    assert input_ops[0].get_module() == functional_model.layers[1]
-    assert isinstance(input_ops[0].get_module(), type(original_model.layers[1].two_convs.conv))
-
-    output_ops = get_all_output_ops(connected_graph)
-    assert len(output_ops) == 1
-    assert output_ops[0].get_module() == functional_model.layers[-1]
-    assert isinstance(output_ops[0].get_module(), type(original_model.layers[-1]))
-
-
-
-def test_multi_math_operations_subclass_to_functional():
-    input_layer = tf.keras.Input(shape=(32, 32, 1))
-    output = MultiMathOperations()(input_layer)
-    original_model = tf.keras.Model(inputs=input_layer, outputs=output)
-
-    functional_model = prepare_model(original_model)
-    assert functional_model.count_params() == original_model.count_params()
-    compare_weights(original_model.get_weights(), functional_model.get_weights())
-
-    # NOTE: No testing of the ConnectedGraph since Lambda layers are not supported by the ConnectedGraph at this time.
-
-
-def test_keras_text_classification_example_model_to_functional():
-    original_model, random_input = get_keras_text_classification_example_model_and_data_input()
-    functional_model = prepare_model(original_model)
-    assert functional_model.count_params() == original_model.count_params()
-
-    # NOTE: Since TextClassification Model has the internal layers out of order compared to the call method,
-    # the weights are not in the order of what the actual architecture is (this is a Keras design).
-    # Therefore, we get the original model's weights and sort them in the order of the actual
-    # architecture and use those weights to compare to the functional model's weights.
-    model_weights_in_correct_order = _get_original_models_weights_in_functional_model_order(
-        original_model, functional_model, class_names={"token_and_position_embedding", "transformer_block"})
-    compare_weights(model_weights_in_correct_order, functional_model.get_weights())
-    verify_functional_model(functional_model,
-                            original_model,
-                            random_input,
-                            number_of_layers_in_model=21)
-
-    # NOTE: No testing of the ConnectedGraph since Embedding, MultiHeadAttention, and Lambda layers are not
-    # supported by the ConnectedGraph at this time.
-
-def test_non_nested_layered_model():
-    original_model = conv_functional()
-    random_input = np.random.rand(32, *original_model.input_shape[1:])
-    orig_output = original_model(random_input)
-
-    functional_model = prepare_model(original_model)
-
-    functional_model_output = functional_model(random_input)
-    model_weights_in_correct_order = _get_original_models_weights_in_functional_model_order(
-        original_model, functional_model, class_names=set())
-
-    compare_weights(model_weights_in_correct_order, functional_model.get_weights())
-    np.testing.assert_array_equal(orig_output.numpy(), functional_model_output.numpy())
-
-def test_multi_output():
-    class TestMultiOut(tf.keras.layers.Layer):
-        def __init__(self, **kwargs) -> None:
-            super().__init__(name='test_multi_out')
-
-        def call(self, inputs, **kwargs):
-            out1 = inputs * 2.0
-            out2 = inputs * 3.0
-            return [out1, out2]
-
-    encoder_input = tf.keras.Input(shape=(28, 28, 1), name="img")
-    x = tf.keras.layers.Conv2D(16, 3, activation="relu")(encoder_input)
-    out = TestMultiOut()(x)
-
-    original_model = tf.keras.Model(inputs=encoder_input, outputs=out)
-    random_input = np.random.rand(1, *original_model.input_shape[1:])
-    _ = original_model(random_input)
-
-    functional_model = prepare_model(original_model)
-    model_weights_in_correct_order = _get_original_models_weights_in_functional_model_order(
-        original_model, functional_model, class_names={'test_multi_out', 'functional'})
-    compare_weights(model_weights_in_correct_order, functional_model.get_weights())
-
-    verify_functional_model(functional_model,
-                            original_model,
-                            random_input,
-                            number_of_layers_in_model=4)
-
-def test_multi_output_only_lambda():
-    encoder_input = tf.keras.Input(shape=(28, 28, 1), name="img")
-    x = tf.keras.layers.Conv2D(16, 3, activation="relu")(encoder_input)
-    x0, x1 = tf.split(x, [8, 8], 3)
-    x1a = x1 * 2.0
-    out = tf.concat([x0, x1a], 3)
-
-    original_model = tf.keras.Model(inputs=encoder_input, outputs=out)
-    random_input = np.random.rand(1, *original_model.input_shape[1:])
-    orig_output = original_model(random_input)
-
-    functional_model = prepare_model(original_model)
-    ConnectedGraph(functional_model)
-    functional_model_output = functional_model(random_input)
-    model_weights_in_correct_order = _get_original_models_weights_in_functional_model_order(
-        original_model, functional_model, class_names=set())
-
-    compare_weights(model_weights_in_correct_order, functional_model.get_weights())
-    np.testing.assert_array_equal(orig_output.numpy(), functional_model_output.numpy())
+        # NOTE: Verification of the model occurs in the prepare_model function.
+        _ = prepare_model(original_model, input_layer=tf.keras.Input(shape=input_shape))
diff --git a/TrainingExtensions/tensorflow/test/python/eager/test_models_keras.py b/TrainingExtensions/tensorflow/test/python/eager/test_models_keras.py
index a941e022b7..6d40aaa877 100644
--- a/TrainingExtensions/tensorflow/test/python/eager/test_models_keras.py
+++ b/TrainingExtensions/tensorflow/test/python/eager/test_models_keras.py
@@ -245,11 +245,130 @@ def transposed_conv_model():
     """
 
     inp = tf.keras.Input((4, 4, 10))
-    x   = tf.keras.layers.Conv2DTranspose(10, 3)(inp)
-    x   = tf.keras.layers.BatchNormalization()(x)
+    x = tf.keras.layers.Conv2DTranspose(10, 3)(inp)
+    x = tf.keras.layers.BatchNormalization()(x)
 
-    x   = tf.keras.layers.ReLU()(x)
-    x   = tf.keras.layers.Conv2DTranspose(10, 3)(x)
-    x   = tf.keras.layers.BatchNormalization()(x)
+    x = tf.keras.layers.ReLU()(x)
+    x = tf.keras.layers.Conv2DTranspose(10, 3)(x)
+    x = tf.keras.layers.BatchNormalization()(x)
 
     return tf.keras.Model(inputs=[inp], outputs=[x])
+
+
+def resnet_34():
+    """
+    ResNet34 Model
+
+    Taken from here https://github.com/safwankdb/ResNet34-TF2/blob/master/model.py
+    """
+
+    class ResBlock(tf.keras.Model):
+        def __init__(self, channels, stride=1, name=''):
+            super(ResBlock, self).__init__(name=name)
+            self.flag = (stride != 1)
+            self.channels = channels
+            self.stride = stride
+            self.conv1 = tf.keras.layers.Conv2D(channels, 3, stride, padding='same')
+            self.bn1 = tf.keras.layers.BatchNormalization()
+            self.conv2 = tf.keras.layers.Conv2D(channels, 3, padding='same')
+            self.bn2 = tf.keras.layers.BatchNormalization()
+            self.relu = tf.keras.layers.ReLU()
+            self.relu_2 = tf.keras.layers.ReLU()
+            if self.flag:
+                self.bn3 = tf.keras.layers.BatchNormalization()
+                self.conv3 = tf.keras.layers.Conv2D(channels, 1, stride)
+
+        def get_config(self):
+            config = super().get_config()
+            config.update({
+                'channels': self.channels,
+                'stride': self.stride,
+                'name': self.name
+            })
+            return config
+
+        def call(self, x):
+            x1 = self.conv1(x)
+            x1 = self.bn1(x1)
+            x1 = self.relu(x1)
+            x1 = self.conv2(x1)
+            x1 = self.bn2(x1)
+            if self.flag:
+                x = self.conv3(x)
+                x = self.bn3(x)
+            x1 = tf.keras.layers.add([x, x1])
+            x1 = self.relu_2(x1)
+            return x1
+
+    class ResNet34(tf.keras.Model):
+        def __init__(self):
+            super(ResNet34, self).__init__(name='ResNet34')
+            self.conv1 = tf.keras.layers.Conv2D(64, 7, 2, padding='same')
+            self.bn = tf.keras.layers.BatchNormalization()
+            self.relu = tf.keras.layers.ReLU()
+            self.mp1 = tf.keras.layers.MaxPooling2D(3, 2)
+
+            self.conv2_1 = ResBlock(64, name='ResBlock')
+            self.conv2_2 = ResBlock(64, name='ResBlock_1')
+            self.conv2_3 = ResBlock(64, name='ResBlock_2')
+
+            self.conv3_1 = ResBlock(128, 2, name='ResBlock_3')
+            self.conv3_2 = ResBlock(128, name='ResBlock_4')
+            self.conv3_3 = ResBlock(128, name='ResBlock_5')
+            self.conv3_4 = ResBlock(128, name='ResBlock_6')
+
+            self.conv4_1 = ResBlock(256, 2, name='ResBlock_7')
+            self.conv4_2 = ResBlock(256, name='ResBlock_8')
+            self.conv4_3 = ResBlock(256, name='ResBlock_9')
+            self.conv4_4 = ResBlock(256, name='ResBlock_10')
+            self.conv4_5 = ResBlock(256, name='ResBlock_11')
+            self.conv4_6 = ResBlock(256, name='ResBlock_12')
+
+            self.conv5_1 = ResBlock(512, 2, name='ResBlock_13')
+            self.conv5_2 = ResBlock(512, name='ResBlock_14')
+            self.conv5_3 = ResBlock(512, name='ResBlock_15')
+
+            self.pool = tf.keras.layers.GlobalAveragePooling2D()
+            self.fc1 = tf.keras.layers.Dense(512, activation='relu')
+            self.dp1 = tf.keras.layers.Dropout(0.5)
+            self.fc2 = tf.keras.layers.Dense(512, activation='relu')
+            self.dp2 = tf.keras.layers.Dropout(0.5)
+            self.fc3 = tf.keras.layers.Dense(64)
+
+        def call(self, x):
+            x = self.conv1(x)
+            x = self.bn(x)
+            x = self.relu(x)
+            x = self.mp1(x)
+
+            x = self.conv2_1(x)
+            x = self.conv2_2(x)
+            x = self.conv2_3(x)
+
+            x = self.conv3_1(x)
+            x = self.conv3_2(x)
+            x = self.conv3_3(x)
+            x = self.conv3_4(x)
+
+            x = self.conv4_1(x)
+            x = self.conv4_2(x)
+            x = self.conv4_3(x)
+            x = self.conv4_4(x)
+            x = self.conv4_5(x)
+            x = self.conv4_6(x)
+
+            x = self.conv5_1(x)
+            x = self.conv5_2(x)
+            x = self.conv5_3(x)
+
+            x = self.pool(x)
+            x = self.fc1(x)
+            x = self.dp1(x)
+            x = self.fc2(x)
+            x = self.dp2(x)
+            x = self.fc3(x)
+            return x
+
+    model = ResNet34()
+    model.build(input_shape=(1, 480, 480, 3))
+    return model