Raise error when a Keras model has a custom layer (#748)

Co-authored-by: reuvenp <reuvenp@altair-semi.com>

model_compression_toolkit/core/keras/custom_layer_validation.py

@@ -0,0 +1,31 @@
+# Copyright 2023 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import inspect
+
+def is_keras_custom_layer(layer_class: type) -> bool:
+    """
+    Check whether a layer class is from keras/tensorflow modules.
+
+    Args:
+        layer_class: Layer class to check its root module.
+
+    Returns:
+        Whether a layer class is from keras/tensorflow modules or not.
+
+    """
+    # Get the root module name the layer is from:
+    root_module_name = inspect.getmodule(layer_class).__name__.split('.')[0]
+    return root_module_name not in ['keras', 'tensorflow']

model_compression_toolkit/core/keras/reader/node_builder.py

@@ -17,6 +17,9 @@
 import tensorflow as tf
 from packaging import version
 
+from model_compression_toolkit.core.keras.custom_layer_validation import is_keras_custom_layer
+from model_compression_toolkit.logger import Logger
+
 if version.parse(tf.__version__) < version.parse("2.6"):
     from tensorflow.python.keras.layers.core import TFOpLambda, SlicingOpLambda
     from tensorflow.python.keras.engine.keras_tensor import KerasTensor
@@ -55,6 +58,12 @@ def build_node(node: KerasNode,
     op_call_args = node.call_args
     op_call_kwargs = node.call_kwargs
     layer_class = type(keras_layer)  # class path to instantiating it in back2framework.
+
+    # Validate the layer is not a custom layer
+    if is_keras_custom_layer(layer_class):
+        Logger.error(f'MCT does not support optimizing Keras custom layers, but found layer of type {layer_class}. '
+                     f'Please file a feature request or an issue if you believe this is an issue.')
+
     weights = {v.name: v.numpy() for v in keras_layer.weights}  # layer's weights
 
     # If it's a node representing a reused layer, several nodes will contain the same layer instance.

tests/keras_tests/function_tests/test_unsupported_custom_layer.py

@@ -0,0 +1,53 @@
+# Copyright 2023 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+import unittest
+
+import numpy as np
+import tensorflow as tf
+
+import model_compression_toolkit as mct
+
+keras = tf.keras
+layers = keras.layers
+
+
+class CustomFC(keras.layers.Layer):
+
+    def __init__(self, units=32, input_dim=3):
+        super().__init__()
+        self.w = self.add_weight(
+            shape=(input_dim, units), initializer="random_normal", trainable=True
+        )
+        self.b = self.add_weight(shape=(units,), initializer="zeros", trainable=True)
+
+    def call(self, inputs):
+        return tf.matmul(inputs, self.w) + self.b
+
+
+class TestUnsupportedCustomLayer(unittest.TestCase):
+
+    def test_raised_error_with_custom_layer(self):
+        inputs = layers.Input(shape=(3, 3, 3))
+        x = CustomFC()(inputs)
+        model = keras.Model(inputs=inputs, outputs=x)
+
+        expected_error = f'MCT does not support optimizing Keras custom layers, but found layer of type <class ' \
+                         f'\'test_unsupported_custom_layer.CustomFC\'>. Please file a feature request or an issue if ' \
+                         f'you believe this is an issue.'
+
+        with self.assertRaises(Exception) as e:
+            mct.ptq.keras_post_training_quantization_experimental(model,
+                                                                  lambda _: [np.random.randn(1, 3, 3, 3)])
+        self.assertEqual(expected_error, str(e.exception))