Fix activation bias correction with preserving node (#1449)

* Fix activation bias correction with preserving node

Author: lapid92

model_compression_toolkit/core/common/mixed_precision/bit_width_setter.py

@@ -76,7 +76,8 @@ def set_bit_widths(mixed_precision_enable: bool,
         for n in graph.nodes:
             assert len(n.candidates_quantization_cfg) == 1
             n.final_weights_quantization_cfg = copy.deepcopy(n.candidates_quantization_cfg[0].weights_quantization_cfg)
-            n.final_activation_quantization_cfg = copy.deepcopy(n.candidates_quantization_cfg[0].activation_quantization_cfg)
+            if not n.is_quantization_preserving():
+                n.final_activation_quantization_cfg = copy.deepcopy(n.candidates_quantization_cfg[0].activation_quantization_cfg)
 
     return graph
 

model_compression_toolkit/core/common/statistics_correction/compute_activation_bias_correction_of_graph.py

@@ -42,13 +42,9 @@ def get_previous_node_with_activation_quantization(linear_node: BaseNode,
 
     prev_node = prev_nodes[0]
 
-    activation_quantization_config = prev_node.final_activation_quantization_cfg
+    prev_quant_node = graph.retrieve_preserved_quantization_node(prev_node)
 
-    # Search for node with activation quantization
-    if activation_quantization_config.enable_activation_quantization:
-        return prev_node
-    else:
-        return get_previous_node_with_activation_quantization(prev_node, graph)
+    return prev_quant_node if prev_quant_node.is_activation_quantization_enabled() else None
 
 
 def calculate_bin_centers(bin_edges: np.ndarray) -> np.ndarray:

tests/keras_tests/feature_networks_tests/feature_networks/activation_bias_correction_test.py

@@ -20,6 +20,7 @@
 import numpy as np
 
 from tests.keras_tests.utils import get_layers_from_model_by_type
+from tests.common_tests.helpers.tpcs_for_tests.v4.tpc import get_tpc
 
 keras = tf.keras
 layers = keras.layers
@@ -52,6 +53,10 @@ def get_quantization_config(self):
                                   activation_bias_correction=True,
                                   activation_bias_correction_threshold=self.activation_bias_correction_threshold)
 
+    def get_tpc(self):
+        tpc = get_tpc()
+        return tpc
+
     def create_networks(self):
         inputs = layers.Input(shape=self.get_input_shapes()[0][1:])
         x = self.prev_layer(inputs)

tests/pytorch_tests/model_tests/feature_models/activation_bias_correction_test.py

@@ -106,6 +106,10 @@ def get_quantization_config(self):
     def create_networks(self):
         return self.model
 
+    def get_tpc(self):
+        from tests.common_tests.helpers.tpcs_for_tests.v4.tpc import get_tpc
+        return get_tpc()
+
     def compare(self, quantized_model, float_model, input_x=None, quantization_info=None):
         bias = float_model.linear_layer.bias.cpu().detach().numpy()
         bias_after_activation_bias_correction = quantized_model.linear_layer.layer.bias.cpu().detach().numpy()