correcting accrding to the feedback comments

gouda-youichi · gouda-youichi · commit b07fc5d88170 · 2025-03-12T10:29:43.000+09:00
diff --git a/model_compression_toolkit/core/common/quantization/bit_width_config.py b/model_compression_toolkit/core/common/quantization/bit_width_config.py
@@ -33,6 +33,54 @@ class ManualBitWidthSelection:
     filter: BaseNodeMatcher
     bit_width: int
 
+@dataclass
+class ManualWeightsBitWidthSelection(ManualBitWidthSelection):
+    """
+   Class to encapsulate the manual weights bit width selection configuration for a specific filter.
+
+   Attributes:
+       filter (BaseNodeMatcher): The filter used to select nodes for bit width manipulation.
+       bit_width (int): The bit width to be applied to the selected nodes.
+   """
+    val: int
+
+
+def _set_manual_bit_width(filters: Union[List[BaseNodeMatcher], BaseNodeMatcher],
+                         bit_widths: Union[List[int], int]):
+
+    filters = [filters] if not isinstance(filters, list) else filters
+    bit_widths = [bit_widths] if not isinstance(bit_widths, list) else bit_widths
+    if len(bit_widths) > 1 and len(bit_widths) != len(filters):
+        Logger.critical(f"Configuration Error: The number of provided bit_width values {len(bit_widths)} "
+                        f"must match the number of filters {len(filters)}, or a single bit_width value "
+                        f"should be provided for all filters.")
+    elif len(bit_widths) == 1 and len(filters) > 1:
+        bit_widths = [bit_widths[0] for f in filters]
+
+    return 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.")
+            unit_nodes_to_change_bit_width.update({n: manual_bit_width_selection.bit_width})
+
+    return unit_nodes_to_change_bit_width
+
+@dataclass
+class NodesToChangeBitWidth:
+    activation_nodes_to_change_bit_width: Dict
+    weights_nodes_to_change_bit_width: Dict
 
 @dataclass
 class BitWidthConfig:
@@ -49,47 +97,33 @@ def set_manual_activation_bit_width(self,
                                         filters: Union[List[BaseNodeMatcher], BaseNodeMatcher],
                                         bit_widths: Union[List[int], int]):
         """
-        Add a manual bit-width selection to the configuration.
+        Add a manual bit-width selection for activation to the configuration.
 
         Args:
-            filter (Union[List[BaseNodeMatcher], BaseNodeMatcher]): The filters used to select nodes for bit-width manipulation.
-            bit_width (Union[List[int], int]): The bit widths to be applied to the selected nodes.
+            filters (Union[List[BaseNodeMatcher], BaseNodeMatcher]): The filters used to select nodes for bit-width manipulation.
+            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
         """
-        filters = [filters] if not isinstance(filters, list) else filters
-        bit_widths = [bit_widths] if not isinstance(bit_widths, list) else bit_widths
-        if len(bit_widths) > 1 and len(bit_widths) != len(filters):
-            Logger.critical(f"Configuration Error: The number of provided bit_width values {len(bit_widths)} "
-                            f"must match the number of filters {len(filters)}, or a single bit_width value "
-                            f"should be provided for all filters.")
-        elif len(bit_widths) == 1 and len(filters) > 1:
-            bit_widths = [bit_widths[0] for f in filters]
+        bit_widths, filters = _set_manual_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)]
 
     def set_manual_weights_bit_width(self,
                                         filters: Union[List[BaseNodeMatcher], BaseNodeMatcher],
                                         bit_widths: Union[List[int], int]):
         """
-        Add a manual bit-width selection to the configuration.
+        Add a manual bit-width selection for weights to the configuration.
 
         Args:
-            filter (Union[List[BaseNodeMatcher], BaseNodeMatcher]): The filters used to select nodes for bit-width manipulation.
-            bit_width (Union[List[int], int]): The bit widths to be applied to the selected nodes.
+            filters (Union[List[BaseNodeMatcher], BaseNodeMatcher]): The filters used to select nodes for bit-width manipulation.
+            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
         """
-        filters = [filters] if not isinstance(filters, list) else filters
-        bit_widths = [bit_widths] if not isinstance(bit_widths, list) else bit_widths
-        if len(bit_widths) > 1 and len(bit_widths) != len(filters):
-            Logger.critical(f"Configuration Error: The number of provided bit_width values {len(bit_widths)} "
-                            f"must match the number of filters {len(filters)}, or a single bit_width value "
-                            f"should be provided for all filters.")
-        elif len(bit_widths) == 1 and len(filters) > 1:
-            bit_widths = [bit_widths[0] for f in filters]
+        bit_widths, filters = _set_manual_bit_width(filters, bit_widths)
         for bit_width, filter in zip (bit_widths, filters):
             self.manual_weights_bit_width_selection_list += [ManualBitWidthSelection(filter, bit_width)]
 
-    def get_nodes_to_manipulate_bit_widths(self, graph: Graph) -> Dict:
+    def get_nodes_to_manipulate_bit_widths(self, graph: Graph) -> NodesToChangeBitWidth:
         """
         Retrieve nodes from the graph that need their bit-widths changed according to the manual bit-width selections.
 
@@ -99,25 +133,9 @@ def get_nodes_to_manipulate_bit_widths(self, graph: Graph) -> Dict:
         Returns:
             Dict: A dictionary mapping nodes to their new bit-widths.
         """
-        def make_nodes_to_change_bit_width(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.")
-                    unit_nodes_to_change_bit_width.update({n: manual_bit_width_selection.bit_width})
-
-            return unit_nodes_to_change_bit_width
-
-        activation_nodes_to_change_bit_width = make_nodes_to_change_bit_width(self.manual_activation_bit_width_selection_list)
-        weights_nodes_to_change_bit_width = make_nodes_to_change_bit_width(self.manual_weights_bit_width_selection_list)
-
-        nodes_to_change_bit_width = {ACTIVATION_ATTRIBUTE: activation_nodes_to_change_bit_width, WEIGHTS_ATTRIBUTE: weights_nodes_to_change_bit_width}
+        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)
+
+        #nodes_to_change_bit_width = {ACTIVATION_ATTRIBUTE: activation_nodes_to_change_bit_width, WEIGHTS_ATTRIBUTE: weights_nodes_to_change_bit_width}
+        nodes_to_change_bit_width = NodesToChangeBitWidth(activation_nodes_to_change_bit_width, weights_nodes_to_change_bit_width)
         return nodes_to_change_bit_width
diff --git a/model_compression_toolkit/core/common/quantization/set_node_quantization_config.py b/model_compression_toolkit/core/common/quantization/set_node_quantization_config.py
@@ -20,7 +20,7 @@
 from mct_quantizers.common.constants import WEIGHTS_N_BITS, ACTIVATION_N_BITS
 from model_compression_toolkit.constants import WEIGHTS_ATTRIBUTE, ACTIVATION_ATTRIBUTE
 from model_compression_toolkit.core.common import BaseNode
-from model_compression_toolkit.core.common.quantization.bit_width_config import BitWidthConfig
+from model_compression_toolkit.core.common.quantization.bit_width_config import BitWidthConfig, NodesToChangeBitWidth
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.core.common.framework_info import FrameworkInfo
 from model_compression_toolkit.core.common.graph.base_graph import Graph
@@ -67,11 +67,13 @@ def set_quantization_configuration_to_graph(graph: Graph,
         Logger.warning("Using the HMSE error method for weights quantization parameters search. "
                        "Note: This method may significantly increase runtime during the parameter search process.")
 
-    nodes_to_manipulate_bit_widths = {} if bit_width_config is None else bit_width_config.get_nodes_to_manipulate_bit_widths(graph)
+    nodes_to_manipulate_bit_widths = NodesToChangeBitWidth() if bit_width_config is None else bit_width_config.get_nodes_to_manipulate_bit_widths(graph)
+    print('nodes_to_manipulate_bit_widths', nodes_to_manipulate_bit_widths)
+    print('zzz', nodes_to_manipulate_bit_widths.activation_nodes_to_change_bit_width)
 
     for n in graph.nodes:
-        manual_bit_width_override = {ACTIVATION_ATTRIBUTE: nodes_to_manipulate_bit_widths.get(ACTIVATION_ATTRIBUTE).get(n),
-                                     WEIGHTS_ATTRIBUTE: nodes_to_manipulate_bit_widths.get(WEIGHTS_ATTRIBUTE).get(n)}
+        manual_bit_width_override = {ACTIVATION_ATTRIBUTE: nodes_to_manipulate_bit_widths.activation_nodes_to_change_bit_width.get(n),
+                                     WEIGHTS_ATTRIBUTE: nodes_to_manipulate_bit_widths.weights_nodes_to_change_bit_width.get(n)}
         set_quantization_configs_to_node(node=n,
                                          graph=graph,
                                          quant_config=quant_config,
diff --git a/tests_pytest/common_tests/core/common/quantization/test_manual_weights_bitwidth_selection.py b/tests_pytest/common_tests/core/common/quantization/test_manual_weights_bitwidth_selection.py
@@ -179,15 +179,15 @@ def test_BitWidthConfig_getter(self, inputs, expected):
         get_manual_bit_dict = manual_bit_cfg.get_nodes_to_manipulate_bit_widths(graph)
 
         if activation[0] is not None:
-            for idx, (key, val) in enumerate(get_manual_bit_dict["activation"].items()):
+            for idx, (key, val) in enumerate(get_manual_bit_dict.activation_nodes_to_change_bit_width.items()):
                 assert str(key) == list(activation_expected.keys())[idx]
                 assert val == list(activation_expected.values())[idx]
         else:
-            assert get_manual_bit_dict["activation"] == activation_expected
+            assert get_manual_bit_dict.activation_nodes_to_change_bit_width == activation_expected
 
         if weights[0] is not None:
-            for idx, (key, val) in enumerate(get_manual_bit_dict["weights"].items()):
+            for idx, (key, val) in enumerate(get_manual_bit_dict.weights_nodes_to_change_bit_width.items()):
                 assert str(key) == list(weights_expected.keys())[idx]
                 assert val == list(weights_expected.values())[idx]
         else:
-            assert get_manual_bit_dict["weights"] == weights_expected
+            assert get_manual_bit_dict.weights_nodes_to_change_bit_width == weights_expected
