Remove activation quantizers from Keras wrapper (#25)

* Remove activation quantizers from Keras wrapper

* Remove call to _set_activation_vars in Keras wrapper

---------

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

Author: reuvenperetz

mct_quantizers/keras/quantize_wrapper.py

@@ -15,7 +15,7 @@
 from typing import Dict, List, Any, Tuple
 
 from mct_quantizers.common.base_inferable_quantizer import BaseInferableQuantizer
-from mct_quantizers.common.constants import FOUND_TF, ACTIVATION_QUANTIZERS, WEIGHTS_QUANTIZERS, STEPS, LAYER, TRAINING
+from mct_quantizers.common.constants import FOUND_TF, WEIGHTS_QUANTIZERS, STEPS, LAYER, TRAINING
 from mct_quantizers.logger import Logger
 from mct_quantizers.common.get_all_subclasses import get_all_subclasses
 
@@ -54,20 +54,17 @@ class KerasQuantizationWrapper(tf.keras.layers.Wrapper):
         def __init__(self,
                      layer,
                      weights_quantizers: Dict[str, BaseInferableQuantizer] = None,
-                     activation_quantizers: List[BaseInferableQuantizer] = None,
                      **kwargs):
             """
             Keras Quantization Wrapper takes a keras layer and quantizers and infer a quantized layer.
 
             Args:
                 layer: A keras layer.
                 weights_quantizers: A dictionary between a weight's name to its quantizer.
-                activation_quantizers: A list of activations quantization, one for each layer output.
             """
             super(KerasQuantizationWrapper, self).__init__(layer, **kwargs)
             self._track_trackable(layer, name='layer')
             self.weights_quantizers = weights_quantizers if weights_quantizers is not None else dict()
-            self.activation_quantizers = activation_quantizers if activation_quantizers is not None else list()
 
         def add_weights_quantizer(self, param_name: str, quantizer: BaseInferableQuantizer):
             """
@@ -82,15 +79,6 @@ def add_weights_quantizer(self, param_name: str, quantizer: BaseInferableQuantiz
             """
             self.weights_quantizers.update({param_name: quantizer})
 
-        @property
-        def is_activation_quantization(self) -> bool:
-            """
-            This function check activation quantizer exists in wrapper.
-            Returns: a boolean if activation quantizer exists
-
-            """
-            return self.num_activation_quantizers > 0
-
         @property
         def is_weights_quantization(self) -> bool:
             """
@@ -108,22 +96,13 @@ def num_weights_quantizers(self) -> int:
             """
             return len(self.weights_quantizers)
 
-        @property
-        def num_activation_quantizers(self) -> int:
-            """
-            Returns: number of activations quantizers
-            """
-            return len(self.activation_quantizers)
-
         def get_config(self):
             """
             Returns: Configuration of KerasQuantizationWrapper.
 
             """
             base_config = super(KerasQuantizationWrapper, self).get_config()
-            config = {
-                ACTIVATION_QUANTIZERS: [keras.utils.serialize_keras_object(act) for act in self.activation_quantizers],
-                WEIGHTS_QUANTIZERS: {k: keras.utils.serialize_keras_object(v) for k, v in self.weights_quantizers.items()}}
+            config = {WEIGHTS_QUANTIZERS: {k: keras.utils.serialize_keras_object(v) for k, v in self.weights_quantizers.items()}}
             return dict(list(base_config.items()) + list(config.items()))
 
         def _set_weights_vars(self, is_training: bool = True):
@@ -143,17 +122,6 @@ def _set_weights_vars(self, is_training: bool = True):
                 self._weights_vars.append((name, weight, quantizer))
                 self._trainable_weights.append(weight) # Must when inherit from tf.keras.layers.Wrapper in tf2.10 and below
 
-        def _set_activations_vars(self):
-            """
-            This function sets activations quantizers vars to the layer
-
-            Returns: None
-            """
-            self._activation_vars = []
-            for i, quantizer in enumerate(self.activation_quantizers):
-                quantizer.initialize_quantization(None, self.layer.name + f'/out{i}', self)
-                self._activation_vars.append(quantizer)
-
         @classmethod
         def from_config(cls, config):
             """
@@ -167,14 +135,11 @@ def from_config(cls, config):
             config = config.copy()
             qi_inferable_custom_objects = {subclass.__name__: subclass for subclass in
                                            get_all_subclasses(BaseKerasInferableQuantizer)}
-            activation_quantizers = [keras.utils.deserialize_keras_object(act,
-                                                                          module_objects=globals(),
-                                                                          custom_objects=None) for act in config.pop(ACTIVATION_QUANTIZERS)]
             weights_quantizers = {k: keras.utils.deserialize_keras_object(v,
                                                                           module_objects=globals(),
                                                                           custom_objects=qi_inferable_custom_objects) for k, v in config.pop(WEIGHTS_QUANTIZERS).items()}
             layer = tf.keras.layers.deserialize(config.pop(LAYER))
-            return cls(layer=layer, weights_quantizers=weights_quantizers, activation_quantizers=activation_quantizers, **config)
+            return cls(layer=layer, weights_quantizers=weights_quantizers, **config)
 
         def build(self, input_shape):
             """
@@ -194,7 +159,6 @@ def build(self, input_shape):
                 trainable=False)
 
             self._set_weights_vars()
-            self._set_activations_vars()
 
         def set_quantize_weights(self, quantized_weights: dict):
             """
@@ -254,29 +218,6 @@ def call(self, inputs, training=None, **kwargs):
             else:
                 outputs = self.layer.call(inputs, **kwargs)
 
-            # Quantize all activations if quantizers exist.
-            if self.is_activation_quantization:
-                num_outputs = len(outputs) if isinstance(outputs, (list, tuple)) else 1
-                if self.num_activation_quantizers != num_outputs:
-                    Logger.error('Quantization wrapper output quantization error: '
-                                 f'number of outputs and quantizers mismatch ({num_outputs}!='
-                                 f'{self.num_activation_quantizers}')
-                if num_outputs == 1:
-                    outputs = [outputs]
-
-                _outputs = []
-                for _output, act_quant in zip(outputs, self.activation_quantizers):
-                    activation_quantizer_args_spec = tf_inspect.getfullargspec(act_quant.__call__).args
-                    if TRAINING in activation_quantizer_args_spec:
-                        _outputs.append(utils.smart_cond(
-                            training,
-                            _make_quantizer_fn(act_quant, _output, True),
-                            _make_quantizer_fn(act_quant, _output, False)))
-                    else:
-                        # Keras activation inferable quantizer.
-                        _outputs.append(act_quant(_output))
-                outputs = _outputs[0] if num_outputs == 1 else _outputs
-
             return outputs
 
         def convert_to_inferable_quantizers(self):
@@ -286,14 +227,6 @@ def convert_to_inferable_quantizers(self):
             Returns:
                 None
             """
-            # Activations quantizers
-            inferable_activation_quantizers = []
-            if self.is_activation_quantization:
-                for quantizer in self.activation_quantizers:
-                    if hasattr(quantizer, 'convert2inferable') and callable(quantizer.convert2inferable):
-                        inferable_activation_quantizers.append(quantizer.convert2inferable())
-                self.activation_quantizers = inferable_activation_quantizers
-
             # Weight quantizers
             inferable_weight_quantizers = {}
             if self.is_weights_quantization:
@@ -310,7 +243,7 @@ def convert_to_inferable_quantizers(self):
                 layer_weights_list.append(getattr(self.layer, weight_attr)) # quantized weights
             layer_weights_list.extend(self.layer.get_weights()) # non quantized weights
             inferable_quantizers_wrapper.layer.set_weights(layer_weights_list)
-            inferable_quantizers_wrapper._set_activations_vars()
+
             # The wrapper inference is using the weights of the quantizers so it expectes to create them by running _set_weights_vars
             inferable_quantizers_wrapper._set_weights_vars(False)
             return inferable_quantizers_wrapper
@@ -342,15 +275,13 @@ def get_quantized_weights(self) -> Dict[str, tf.Tensor]:
     class KerasQuantizationWrapper(object):
         def __init__(self,
                      layer,
-                     weights_quantizers: Dict[str, BaseInferableQuantizer] = None,
-                     activation_quantizers: List[BaseInferableQuantizer] = None):
+                     weights_quantizers: Dict[str, BaseInferableQuantizer] = None):
             """
             Keras Quantization Wrapper takes a keras layer and quantizers and infer a quantized layer.
 
             Args:
                 layer: A keras layer.
                 weights_quantizers: A dictionary between a weight's name to its quantizer.
-                activation_quantizers: A list of activations quantization, one for each layer output.
             """
             Logger.critical('Installing tensorflow and tensorflow_model_optimization is mandatory '
                             'when using KerasQuantizationWrapper. '

tests/keras_tests/test_keras_quantization_wrapper.py

@@ -42,20 +42,6 @@ def initialize_quantization(self, tensor_shape, name, layer):
         return {}
 
 
-class ZeroActivationsQuantizer:
-    """
-    A dummy quantizer for test usage - "quantize" the layer's activation to 0
-    """
-
-    def __call__(self,
-                 inputs: tf.Tensor,
-                 training: bool = True) -> tf.Tensor:
-        return inputs * 0
-
-    def initialize_quantization(self, tensor_shape, name, layer):
-        return {}
-
-
 class TestKerasQuantizationWrapper(unittest.TestCase):
 
     def setUp(self):
@@ -86,17 +72,3 @@ def test_weights_quantization_wrapper(self):
         outputs = wrapper.call(call_inputs.astype('float32'))
         self.assertTrue((outputs == conv_layer(call_inputs)).numpy().all())
 
-    def test_activation_quantization_wrapper(self):
-        conv_layer = self.model.layers[1]
-
-        wrapper = KerasQuantizationWrapper(conv_layer, activation_quantizers=[ZeroActivationsQuantizer()])
-
-        # build
-        wrapper.build(self.input_shapes)
-        (act_quantizer) = wrapper._activation_vars[0]
-        self.assertTrue(isinstance(act_quantizer, ZeroActivationsQuantizer))
-
-        # apply the wrapper on inputs
-        call_inputs = self.inputs[0]
-        outputs = wrapper.call(call_inputs.astype('float32'))
-        self.assertTrue((outputs == 0).numpy().all())