diff --git a/model_compression_toolkit/core/keras/custom_layer_validation.py b/model_compression_toolkit/core/keras/custom_layer_validation.py new file mode 100644 index 000000000..0220f0e10 --- /dev/null +++ b/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'] diff --git a/model_compression_toolkit/core/keras/reader/node_builder.py b/model_compression_toolkit/core/keras/reader/node_builder.py index ca9a0bec6..c05dc9d41 100644 --- a/model_compression_toolkit/core/keras/reader/node_builder.py +++ b/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. diff --git a/tests/keras_tests/function_tests/test_unsupported_custom_layer.py b/tests/keras_tests/function_tests/test_unsupported_custom_layer.py new file mode 100644 index 000000000..18704ff3c --- /dev/null +++ b/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 . 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))