fixed PR-FB for manual weights selection bitwidth

Author: gouda-youichi

model_compression_toolkit/core/common/quantization/bit_width_config.py

@@ -20,6 +20,7 @@
 from model_compression_toolkit.core.common.matchers.node_matcher import BaseNodeMatcher
 from model_compression_toolkit.logger import Logger
 
+from model_compression_toolkit.core.common.graph.base_node import WeightAttrT
 
 @dataclass
 class ManualBitWidthSelection:
@@ -40,68 +41,22 @@ class ManualWeightsBitWidthSelection(ManualBitWidthSelection):
 
    Attributes:
        filter (BaseNodeMatcher): The filter used to select nodes for bit width manipulation.
-       attr (str): The attribute used to select nodes for bit width manipulation. [KERNEL_ATTR|BIAS_ATTR]
+       bit_width (int): The bit width to be applied to the selected nodes.
+       attr (str): The filtered node's attributes to apply bit-width manipulation to.
    """
-    attr: str
-
-def _expand_to_list_core(
-        filters: Union[List[BaseNodeMatcher]],
-        vals: Union[List[any], any]):
-    vals = [vals] if not isinstance(vals, list) else vals
-    if len(vals) > 1 and len(vals) != len(filters):
-        Logger.critical(f"Configuration Error: The number of provided bit_width values {len(vals)} "
-                        f"must match the number of filters {len(filters)}, or a single bit_width value "
-                        f"should be provided for all filters.")
-    elif len(vals) == 1 and len(filters) > 1:
-        vals = [vals[0] for f in filters]
-    return vals
-
-def _expand_to_list_filter_and_bit_width(
-        filters: Union[List[BaseNodeMatcher]],
-        bit_widths: Union[List[int], int],
-        attrs: Union[List[str], str] = None):
-    filters = [filters] if not isinstance(filters, list) else filters
-    bit_widths = _expand_to_list_core(filters, bit_widths)
-    attrs = _expand_to_list_core(filters, attrs)
-
-    return attrs, bit_widths, filters
-
-def _make_nodes_to_change_bit_width(graph, manual_bit_width_selection_list):
-    unit_nodes_to_change_bit_width = {}
-    for manual_bit_width_selection in manual_bit_width_selection_list:
-        filtered_nodes = graph.filter(manual_bit_width_selection.filter)
-        if len(filtered_nodes) == 0:
-            Logger.critical(
-                f"Node Filtering Error: No nodes found in the graph for filter {manual_bit_width_selection.filter.__dict__} "
-                f"to change their bit width to {manual_bit_width_selection.bit_width}.")
-        for n in filtered_nodes:
-            # check if a manual configuration exists for this node
-            if n in unit_nodes_to_change_bit_width:
-                Logger.info(
-                    f"Node {n} has an existing manual bit width configuration of {unit_nodes_to_change_bit_width.get(n)}. A new manual configuration request of {manual_bit_width_selection.bit_width} has been received, and the previous value is being overridden.")
-            if isinstance(manual_bit_width_selection_list,  ManualBitWidthSelection):
-                unit_nodes_to_change_bit_width.update({n: manual_bit_width_selection.bit_width})
-            elif isinstance(manual_bit_width_selection_list, ManualWeightsBitWidthSelection):
-                unit_nodes_to_change_bit_width.update({n: [manual_bit_width_selection.bit_width, manual_bit_width_selection.attr]})
-
-
-    return unit_nodes_to_change_bit_width
-
-@dataclass
-class NodesToChangeBitWidth:
-    activation_nodes_to_change_bit_width: Dict
-    weights_nodes_to_change_bit_width: Dict
+    attr: WeightAttrT
 
 @dataclass
 class BitWidthConfig:
     """
     Class to manage manual bit-width configurations.
 
     Attributes:
-        manual_activation_bit_width_selection_list (List[ManualBitWidthSelection]): A list of ManualBitWidthSelection objects defining manual bit-width configurations.
+        manual_activation_bit_width_selection_list (List[ManualBitWidthSelection]): A list of ManualBitWidthSelection objects for activation defining manual bit-width configurations.
+        manual_activation_bit_width_selection_list (List[ManualWeightsBitWidthSelection]): A list of ManualWeightsBitWidthSelection for weights objects defining manual bit-width configurations.
     """
     manual_activation_bit_width_selection_list: List[ManualBitWidthSelection] = field(default_factory=list)
-    manual_weights_bit_width_selection_list: List[ManualBitWidthSelection] = field(default_factory=list)
+    manual_weights_bit_width_selection_list: List[ManualWeightsBitWidthSelection] = field(default_factory=list)
 
     def set_manual_activation_bit_width(self,
                                         filters: Union[List[BaseNodeMatcher], BaseNodeMatcher],
@@ -114,7 +69,7 @@ def set_manual_activation_bit_width(self,
             bit_widths (Union[List[int], int]): The bit widths to be applied to the selected nodes.
             If a single value is given it will be applied to all the filters
         """
-        _, bit_widths, filters = _expand_to_list_filter_and_bit_width(filters, bit_widths)
+        _, bit_widths, filters = self._expand_to_list_filter_and_bit_width(filters, bit_widths)
         for bit_width, filter in zip (bit_widths, filters):
             self.manual_activation_bit_width_selection_list += [ManualBitWidthSelection(filter, bit_width)]
 
@@ -128,15 +83,15 @@ def set_manual_weights_bit_width(self,
 
         Args:
             filters (Union[List[BaseNodeMatcher], BaseNodeMatcher]): The filters used to select nodes for bit-width manipulation.
-            bit_widths (Union[List[int], int]): The bit widths for kernel to be applied to the selected nodes.
-            attrs (Union[List[str], str]): The attributes used to select nodes for bit width manipulation. [KERNEL_ATTR|BIAS_ATTR]
+            bit_widths (Union[List[int], int]): The bit widths for specified by attrs to be applied to the selected nodes.
+            attrs (Union[List[str], str]): The filtered node's attributes to apply bit-width manipulation to.
             If a single value is given it will be applied to all the filters
         """
-        attrs, bit_widths, filters = _expand_to_list_filter_and_bit_width(filters, bit_widths, attrs)
+        attrs, bit_widths, filters = self._expand_to_list_filter_and_bit_width(filters, bit_widths, attrs)
         for attr, bit_width, filter in zip (attrs, bit_widths, filters):
             self.manual_weights_bit_width_selection_list += [ManualWeightsBitWidthSelection(filter, bit_width, attr)]
 
-    def get_nodes_to_manipulate_bit_widths(self, graph: Graph) -> NodesToChangeBitWidth:
+    def get_nodes_to_manipulate_activation_bit_widths(self, graph: Graph) -> Dict:
         """
         Retrieve nodes from the graph that need their bit-widths changed according to the manual bit-width selections.
 
@@ -146,8 +101,68 @@ def get_nodes_to_manipulate_bit_widths(self, graph: Graph) -> NodesToChangeBitWi
         Returns:
             Dict: A dictionary mapping nodes to their new bit-widths.
         """
-        activation_nodes_to_change_bit_width = _make_nodes_to_change_bit_width(graph, self.manual_activation_bit_width_selection_list)
-        weights_nodes_to_change_bit_width = _make_nodes_to_change_bit_width(graph, self.manual_weights_bit_width_selection_list)
+        activation_nodes_to_change_bit_width = self._construct_node_to_new_bit_mapping(graph, self.manual_activation_bit_width_selection_list)
+
+        return activation_nodes_to_change_bit_width
+
+    def get_nodes_to_manipulate_weights_bit_widths(self, graph: Graph) -> Dict:
+        """
+        Retrieve nodes from the graph that need their bit-widths changed according to the manual bit-width selections.
 
-        nodes_to_change_bit_width = NodesToChangeBitWidth(activation_nodes_to_change_bit_width, weights_nodes_to_change_bit_width)
-        return nodes_to_change_bit_width
+        Args:
+            graph (Graph): The graph containing the nodes to be filtered and manipulated.
+
+        Returns:
+            Dict: A dictionary mapping nodes to their new bit-widths.
+        """
+        weights_nodes_to_change_bit_width = self._construct_node_to_new_bit_mapping(graph, self.manual_weights_bit_width_selection_list)
+
+        return weights_nodes_to_change_bit_width
+
+
+
+    def _expand_to_list_core(
+            self,
+            filters: Union[List[BaseNodeMatcher]],
+            vals: Union[List[any], any]):
+        vals = [vals] if not isinstance(vals, list) else vals
+        if len(vals) > 1 and len(vals) != len(filters):
+            Logger.critical(f"Configuration Error: The number of provided bit_width values {len(vals)} "
+                            f"must match the number of filters {len(filters)}, or a single bit_width value "
+                            f"should be provided for all filters.")
+        elif len(vals) == 1 and len(filters) > 1:
+            vals = [vals[0] for f in filters]
+        return vals
+
+    def _expand_to_list_filter_and_bit_width(
+            self,
+            filters: Union[List[BaseNodeMatcher]],
+            bit_widths: Union[List[int], int],
+            attrs: Union[List[str], str] = None):
+        filters = [filters] if not isinstance(filters, list) else filters
+        bit_widths = self._expand_to_list_core(filters, bit_widths)
+        attrs = self._expand_to_list_core(filters, attrs)
+
+        return attrs, bit_widths, filters
+
+    def _construct_node_to_new_bit_mapping(self, graph, manual_bit_width_selection_list):
+        unit_nodes_to_change_bit_width = {}
+        for manual_bit_width_selection in manual_bit_width_selection_list:
+            filtered_nodes = graph.filter(manual_bit_width_selection.filter)
+            if len(filtered_nodes) == 0:
+                Logger.critical(
+                    f"Node Filtering Error: No nodes found in the graph for filter {manual_bit_width_selection.filter.__dict__} "
+                    f"to change their bit width to {manual_bit_width_selection.bit_width}.")
+            for n in filtered_nodes:
+                if n.get_node_weights_attributes() is False:
+                    Logger.critical(f'The requested attribute to change the bit width for {n} is not existing.')
+                # check if a manual configuration exists for this node
+                if n in unit_nodes_to_change_bit_width:
+                    Logger.info(
+                        f"Node {n} has an existing manual bit width configuration of {unit_nodes_to_change_bit_width.get(n)}. A new manual configuration request of {manual_bit_width_selection.bit_width} has been received, and the previous value is being overridden.")
+                if isinstance(manual_bit_width_selection_list,  ManualBitWidthSelection):
+                    unit_nodes_to_change_bit_width.update({n: manual_bit_width_selection.bit_width})
+                elif isinstance(manual_bit_width_selection_list, ManualWeightsBitWidthSelection):
+                    unit_nodes_to_change_bit_width.update({n: [manual_bit_width_selection.bit_width, manual_bit_width_selection.attr]})
+
+        return unit_nodes_to_change_bit_width

tests_pytest/common_tests/core/common/quantization/test_manual_weights_bitwidth_selection.py

@@ -55,9 +55,7 @@ def get_test_graph():
     return graph
 
 class TestBitWidthConfig:
-
-    #######################################################################################################
-    ### test case
+    # test case
     setter_test_input_0 = {"activation": (None, None),
                            "weights":    (None, None, None)}
     setter_test_input_1 = {"activation": (NodeTypeFilter(ReLU), [16]),
@@ -89,7 +87,7 @@ class TestBitWidthConfig:
         (setter_test_input_3, setter_test_expected_3),
         (setter_test_input_4, setter_test_expected_4),
     ])
-    def test_BitWidthConfig_setter(self, inputs, expected):
+    def test_bit_width_config_setter(self, inputs, expected):
 
         def check_param(mb_cfg, exp):
             ### check setting config class (expected ManualBitWidthSelection)
@@ -157,7 +155,6 @@ def check_param_for_weights(mb_cfg, exp):
                 check_param_for_weights(w_mb_cfg, weights_expected[idx])
 
 
-    #######################################################################################################
     ### test case
     ### Note: setter inputs reuse getters test inputs
     getter_test_expected_0 = {"activation":{},
@@ -179,7 +176,7 @@ def check_param_for_weights(mb_cfg, exp):
         (setter_test_input_3, getter_test_expected_3),
         (setter_test_input_4, getter_test_expected_4),
     ])
-    def test_BitWidthConfig_getter(self, inputs, expected):
+    def test_bit_width_config_getter(self, inputs, expected):
 
         graph = get_test_graph()
 
@@ -195,19 +192,20 @@ def test_BitWidthConfig_getter(self, inputs, expected):
         if weights[0] is not None:
             manual_bit_cfg.set_manual_weights_bit_width(weights[0], weights[1], weights[2])
 
-        get_manual_bit_dict = manual_bit_cfg.get_nodes_to_manipulate_bit_widths(graph)
+        get_manual_bit_dict_activation = manual_bit_cfg.get_nodes_to_manipulate_activation_bit_widths(graph)
+        get_manual_bit_dict_weights = manual_bit_cfg.get_nodes_to_manipulate_weights_bit_widths(graph)
 
         if activation[0] is not None:
-            for idx, (key, val) in enumerate(get_manual_bit_dict.activation_nodes_to_change_bit_width.items()):
+            for idx, (key, val) in enumerate(get_manual_bit_dict_activation.items()):
                 assert str(key) == list(activation_expected.keys())[idx]
                 assert val == list(activation_expected.values())[idx]
         else:
-            assert get_manual_bit_dict.activation_nodes_to_change_bit_width == activation_expected
+            assert get_manual_bit_dict_activation == activation_expected
 
         if weights[0] is not None:
-            for idx, (key, val) in enumerate(get_manual_bit_dict.weights_nodes_to_change_bit_width.items()):
+            for idx, (key, val) in enumerate(get_manual_bit_dict_weights.items()):
                 assert str(key) == list(weights_expected.keys())[idx]
                 assert val[0] == list(weights_expected.values())[idx][0]
                 assert val[1] == list(weights_expected.values())[idx][1]
         else:
-            assert get_manual_bit_dict.weights_nodes_to_change_bit_width == weights_expected
+            assert get_manual_bit_dict_weights == weights_expected