align setting fln op config with the new code

Author: irenab

model_compression_toolkit/core/common/graph/base_graph.py

@@ -860,32 +860,36 @@ def _find_intermediate_and_exit_nodes(self, entry_node: BaseNode, fw_impl: Any)
 
         return intermediate_nodes, next_node
 
+    # TODO irena move to load_fqc and clean up tests (currently tests_pytest/common_tests/unit_tests/core/graph/test_base_graph.py)
     def override_fused_node_activation_quantization_candidates(self):
         """
         Override fused node activation quantization candidates for all nodes in fused operations,
         except for the last node in each fused group.
         Update the value of quantization_config with the value of op_quaitization_cfg from FusingInfo.
         """
-        from model_compression_toolkit.core.common.quantization.candidate_node_quantization_config import CandidateNodeQuantizationConfig
-
         nodes_in_fln = self.fusing_info.get_inner_fln_nodes()
         for node in nodes_in_fln:
             fused_node_op_id = self.fusing_info.get_fused_op_id_for_node(node.name)
-            fusiong_op_quaitization_cfg = self.fusing_info.get_fused_op_quantization_config(fused_node_op_id)
-            org_candidate = node.candidates_quantization_cfg[0]
-            if fusiong_op_quaitization_cfg is not None and fusiong_op_quaitization_cfg.enable_activation_quantization:
-                # Set ActivationQuantizationMode to FLN_QUANT and update the value of quantization_config
-                activation_quantization_cfg = NodeActivationQuantizationConfig(qc=org_candidate,
-                                                                               op_cfg=fusiong_op_quaitization_cfg,
-                                                                               activation_quantization_fn=org_candidate.activation_quantization_cfg.activation_quantization_fn,
-                                                                               activation_quantization_params_fn=org_candidate.activation_quantization_cfg.activation_quantization_params_fn)
-                activation_quantization_cfg.quant_mode = ActivationQuantizationMode.FLN_QUANT
-                for qc in node.candidates_quantization_cfg:
-                    qc.activation_quantization_cfg = activation_quantization_cfg
+            fusing_op_quantization_cfg = self.fusing_info.get_fused_op_quantization_config(fused_node_op_id)
+            if fusing_op_quantization_cfg is not None and fusing_op_quantization_cfg.enable_activation_quantization:
+                def update(qc):
+                    qc.activation_quantization_cfg = NodeActivationQuantizationConfig(
+                        fusing_op_quantization_cfg,
+                        qc.activation_quantization_cfg.activation_quantization_fn,
+                        qc.activation_quantization_cfg.activation_quantization_params_fn
+                    )
+                    qc.activation_quantization_cfg.quant_mode = ActivationQuantizationMode.FLN_QUANT
+                node.quantization_cfg.update_all(update)
+                node.quantization_cfg.remove_duplicates()
             else:
-                # Set ActivationQuantizationMode to FLN_NO_QUANT
+                node.quantization_cfg.update_activation_quantization_mode(ActivationQuantizationMode.FLN_NO_QUANT)
+                # Remove duplicate candidates. We cannot compare whole candidates since activation configs might not
+                # be identical, but we do want to treat them as such. So we only check duplication by weight configs.
+                uniq_qcs = []
                 for qc in node.candidates_quantization_cfg:
-                    qc.activation_quantization_cfg.quant_mode = ActivationQuantizationMode.FLN_NO_QUANT
+                    if not any(qc.weights_quantization_cfg == uqc.weights_quantization_cfg for uqc in uniq_qcs):
+                        uniq_qcs.append(qc)
+                node.quantization_cfg.candidates_quantization_cfg = uniq_qcs
 
     def validate(self):
         """

model_compression_toolkit/core/common/quantization/candidate_node_quantization_config.py

@@ -89,6 +89,16 @@ def get_activation_quant_mode(self) -> ActivationQuantizationMode:
         self._validate_consistent_activation_quant_mode()
         return self.base_quantization_cfg.activation_quantization_cfg.quant_mode
 
+    def remove_duplicates(self):
+        """
+        Remove duplicate candidates. First candidate among duplicates is kept, and the order is preserved.
+        """
+        uniq_qcs = []
+        for qc in self.candidates_quantization_cfg:
+            if qc not in uniq_qcs:
+                uniq_qcs.append(qc)
+        self.candidates_quantization_cfg = uniq_qcs
+
     def __post_init__(self, validate=True):
         if validate:
             if not any(self.base_quantization_cfg == qc for qc in self.candidates_quantization_cfg):

model_compression_toolkit/quantization_preparation/load_fqc.py

@@ -139,11 +139,22 @@ def _set_nodes_quantization_configuration(graph: Graph,
     return graph
 
 
-def _set_fusion_info(graph, fqc) -> Graph:
+def _set_fusion_info(graph: Graph, fqc: FrameworkQuantizationCapabilities) -> Graph:
+    """
+
+    Args:
+        graph: graph.
+        fqc: quantization capabilities with attached framework.
+
+    Returns:
+
+    """
     # TODO fix the dict with const keys inside get_fusing_patterns. use named tuple or class
+    # TODO irena instead of storing fusion inside graph (including tpc objects) and then let graph convert tpc op config to
+    #  node config, do it here and only store in graph whatever is relevant after this stage.
     fusing_info = FusingInfoGenerator(fqc.get_fusing_patterns()).generate_fusing_info(graph)
     graph.fusing_info = fusing_info
-    graph.disable_fused_nodes_activation_quantization()
+    graph.override_fused_node_activation_quantization_candidates()
     return graph
 
 

tests_pytest/common_tests/unit_tests/core/graph/test_base_graph.py

@@ -12,10 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+import itertools
 from copy import deepcopy 
 
 import pytest
-from unittest.mock import Mock
+from unittest.mock import Mock, PropertyMock
 
 from mct_quantizers import QuantizationMethod
 from model_compression_toolkit.core.common import Graph
@@ -24,10 +25,13 @@
 from model_compression_toolkit.target_platform_capabilities.schema.mct_current_schema import Signedness
 from tests.common_tests.helpers.generate_test_tpc import generate_test_attr_configs, generate_test_op_qc
 from model_compression_toolkit.core.common.quantization.node_quantization_config import ActivationQuantizationMode, NodeActivationQuantizationConfig
-from model_compression_toolkit.core.common.quantization.candidate_node_quantization_config import CandidateNodeQuantizationConfig
+from model_compression_toolkit.core.common.quantization.candidate_node_quantization_config import \
+    CandidateNodeQuantizationConfig, NodeQuantizationConfig
 from model_compression_toolkit.core.common.quantization.quantization_params_generation.power_of_two_selection import power_of_two_selection_histogram
 from model_compression_toolkit.core.common.quantization.quantization_params_generation.symmetric_selection import symmetric_selection_histogram
 from model_compression_toolkit.core import QuantizationErrorMethod
+from tests_pytest._test_util.graph_builder_utils import build_node, build_nbits_qc
+
 
 def build_mock_fusing_info(nodes, idx):
     """
@@ -55,36 +59,26 @@ def build_mock_fusing_info(nodes, idx):
 
     return fusing_info
 
-def build_mock_node(name, layer_class):
+def build_mock_node(name, layer_class, w_cfgs):
     """
     Creates mock nodes representing a simple neural network structure.
     """
-    node = Mock(spec=BaseNode)
-    node.name = name
-    node.layer_class = layer_class
-
-    activation_quantization_cfg = Mock(spec=NodeActivationQuantizationConfig)
-    activation_quantization_cfg.quant_mode = Mock()
-    activation_quantization_cfg.activation_quantization_fn = symmetric_selection_histogram
-    activation_quantization_cfg.activation_quantization_params_fn = power_of_two_selection_histogram
-   
-    candidate_quantization_config = Mock(spec=CandidateNodeQuantizationConfig)
-    candidate_quantization_config.activation_quantization_cfg = activation_quantization_cfg
-    candidate_quantization_config.activation_error_method = QuantizationErrorMethod.MSE
-    candidate_quantization_config.relu_bound_to_power_of_2 = 0
-    candidate_quantization_config.activation_channel_equalization = False
-    candidate_quantization_config.input_scaling = False
-    candidate_quantization_config.min_threshold = 0
-    candidate_quantization_config.l_p_value = 0
-    candidate_quantization_config.shift_negative_activation_correction = 0
-    candidate_quantization_config.z_threshold = 0
-    candidate_quantization_config.shift_negative_ratio = 0
-    candidate_quantization_config.shift_negative_threshold_recalculation = 0
-    candidate_quantization_config.concat_threshold_update = 0
-    candidate_quantization_config.weights_quantization_cfg = 0
-
-    node.candidates_quantization_cfg = [candidate_quantization_config]
+    node = build_node(name, layer_class=layer_class)
+
+    def eq(self_, other):
+        return self_.activation_n_bits == other.activation_n_bits and self_._quant_mode == other.quant_mode
+    a_cfgs = [Mock(spec=NodeActivationQuantizationConfig,
+                   quant_mode=Mock(),
+                   activation_n_bits=b,
+                   activation_quantization_fn=symmetric_selection_histogram,
+                   activation_quantization_params_fn=power_of_two_selection_histogram,
+                   __eq__=eq) for b in [5, 6]]
+
+    qcs = [CandidateNodeQuantizationConfig(a_cfg, w_cfg) for a_cfg, w_cfg in itertools.product(a_cfgs, w_cfgs)]
 
+    node.quantization_cfg = NodeQuantizationConfig(base_quantization_cfg=qcs[0],
+                                                   candidates_quantization_cfg=qcs,
+                                                   validate=False)
     return node
 
 
@@ -95,14 +89,15 @@ class TestGraph:
         2,
         3,
     ])
-    def test_override_fused_node_activation_quantization_candidates(self, idx):
+    def test_override_fused_node_activation_quantization_candidates(self, idx, patch_fw_info):
         """
         Test the override_fused_node_activation_quantization_candidates function for a graph with multiple nodes and configurations.
         """
         ### Create Test Nodes
         mock_nodes = []
-        mock_nodes.append(build_mock_node(name='conv', layer_class='Conv2d'))
-        mock_nodes.append(build_mock_node(name='fc', layer_class='Linear'))
+        w_cfgs = [Mock(), Mock()]
+        mock_nodes.append(build_mock_node(name='conv', layer_class='Conv2d', w_cfgs=w_cfgs))
+        mock_nodes.append(build_mock_node(name='fc', layer_class='Linear', w_cfgs=w_cfgs[:1]))
 
         ### Create a mock graph
         ### Note: Generate the graph first because fusing_info cannot be set without it.
@@ -120,14 +115,33 @@ def test_override_fused_node_activation_quantization_candidates(self, idx):
         nodes = list(graph.nodes)
 
         if idx == 1:
-            ### Check if the first node ActivationQuantization settings match the expected values
-            assert nodes[0].candidates_quantization_cfg[0].activation_quantization_cfg.quant_mode == ActivationQuantizationMode.FLN_QUANT
-            assert nodes[0].candidates_quantization_cfg[0].activation_quantization_cfg.activation_n_bits == 16
-            assert nodes[0].candidates_quantization_cfg[0].activation_quantization_cfg.signedness == Signedness.AUTO
-            assert nodes[0].candidates_quantization_cfg[0].activation_quantization_cfg.activation_quantization_method == QuantizationMethod.POWER_OF_TWO
-            assert nodes[0].candidates_quantization_cfg[0].activation_quantization_cfg.activation_quantization_params_fn == power_of_two_selection_histogram
+            # Check if the first node ActivationQuantization settings match the expected values
+            # Weight mp configs are preserved, all candidates have the new activation config and duplicates are removed
+            qcs0 = nodes[0].quantization_cfg.candidates_quantization_cfg
+            assert len(qcs0) == 2
+            for i, qc in enumerate(qcs0):
+                assert qc.activation_quantization_cfg.quant_mode == ActivationQuantizationMode.FLN_QUANT
+                assert qc.activation_quantization_cfg.activation_n_bits == 16
+                assert qc.activation_quantization_cfg.signedness == Signedness.AUTO
+                assert qc.activation_quantization_cfg.activation_quantization_method == QuantizationMethod.POWER_OF_TWO
+                assert qc.activation_quantization_cfg.activation_quantization_params_fn == power_of_two_selection_histogram
+                assert qc.weights_quantization_cfg == w_cfgs[i]
+            base_cfg0 = nodes[0].quantization_cfg.base_quantization_cfg
+            assert base_cfg0.activation_quantization_cfg.activation_n_bits == 16
+            assert base_cfg0.activation_quantization_cfg.quant_mode == ActivationQuantizationMode.FLN_QUANT
             ### Check if the second node ActivationQuantization settings match the expected values
-            assert nodes[1].candidates_quantization_cfg[0].activation_quantization_cfg.quant_mode == ActivationQuantizationMode.FLN_NO_QUANT
+            # activations are fln-disabled, duplicates are removed even though orig activation configs differ in nbits
+            qcs1 = nodes[1].quantization_cfg.candidates_quantization_cfg
+            assert len(qcs1) == 1
+            assert qcs1[0].activation_quantization_cfg.quant_mode == ActivationQuantizationMode.FLN_NO_QUANT
+            assert qcs1[0].weights_quantization_cfg == w_cfgs[0]
+            assert (nodes[1].quantization_cfg.base_quantization_cfg.
+                    activation_quantization_cfg.quant_mode == ActivationQuantizationMode.FLN_NO_QUANT)
+
         else:
             ### Check if the first node ActivationQuantization settings match the expected values
-            assert nodes[0].candidates_quantization_cfg[0].activation_quantization_cfg.quant_mode == ActivationQuantizationMode.FLN_NO_QUANT
+            qcs0 = nodes[0].quantization_cfg.candidates_quantization_cfg
+            assert len(qcs0) == 2
+            for i, qc in enumerate(qcs0):
+                assert qc.activation_quantization_cfg.quant_mode == ActivationQuantizationMode.FLN_NO_QUANT
+                assert qc.weights_quantization_cfg == w_cfgs[i]