Raising exception when model with seperable conv layer calls QuantSim before running model preparer

TrainingExtensions/tensorflow/src/python/aimet_tensorflow/keras/quantsim.py

@@ -35,6 +35,7 @@
 #  @@-COPYRIGHT-END-@@
 # =============================================================================
 """ Quantsim for Keras """
+from __future__ import annotations
 
 from dataclasses import dataclass
 import json
@@ -149,6 +150,28 @@ def _validate_model(self):
             _logger.error(error_msg)
             raise NotImplementedError(error_msg)
 
+        sep_conv_found = self.check_separable_conv(self._model_without_wrappers)
+        if sep_conv_found:
+            # Raising an assertion error incase there's SeparableConv2D in the model because in this case we have two sets of weights: Depthwise
+            # and Pointwise. For depthwise kernels, LAST TWO AXIS should be considered and for pointwise kernels LAST AXIS
+            # should be considered, which is not handled here. Running model preparer beforehand will resolve this as there the
+            # SeparableConv2D is splitted into two layers Depthwise and Pointwise seperately.
+            raise AssertionError("SeparableConv2D found in the model. Please run model preparer before calling QuantizationSimModel")
+
+    def check_separable_conv(self, model: tf.keras.models.Model | tf.keras.Sequential) -> bool:
+        """
+        Checks for SeparableConv2D layer in the model
+        :param model: Keras Model
+        :return: Boolean value, True if SeperableConv layer is found else False
+        """
+        for layer in model.layers:
+            if isinstance(layer, tf.keras.Sequential):
+                if self.check_separable_conv(layer):
+                    return True
+            elif isinstance(layer, tf.keras.layers.SeparableConv2D):
+                return True
+        return False
+
     def _get_quantizer_list(self) -> Tuple[List, List, List]:
         """
         Method to provide a list of input, output and parameter quantizers

TrainingExtensions/tensorflow/test/python/eager/test_quantsim_keras.py

@@ -1411,3 +1411,28 @@ def test_multi_output_model():
 
     with tempfile.TemporaryDirectory() as tmp_dir:
         sim.export(tmp_dir, "multi_output_model")
+
+def test_quantsim_with_separable_conv():
+    """Tests Quantsim with seperable conv layer without model preparer """
+
+    tf.keras.backend.clear_session()
+    input_shape = (2, 3, 3, 2)
+    inp = tf.keras.layers.Input(shape=input_shape[1:])
+    test_inp = np.random.rand(*input_shape)
+
+    separable_conv = tf.keras.layers.SeparableConv2D(
+        2,
+        kernel_size=2,
+        depthwise_initializer=tf.initializers.Constant([5.0]),
+        bias_initializer=tf.initializers.Constant([5.0])
+    )(inp)
+    model = tf.keras.models.Model(inputs=[inp], outputs=[separable_conv])
+
+    _ = model(test_inp)
+
+    with pytest.raises(AssertionError) as exc_info:
+        sim = QuantizationSimModel(model, quant_scheme='tf_enhanced', rounding_mode='nearest',
+                                   default_output_bw=8, default_param_bw=8)
+
+    assert str(exc_info.value) == 'SeparableConv2D found in the model. Please run model preparer before calling ' \
+                                  'QuantizationSimModel'