Integrate manual weights bitwidth selection into the MCT core (#1395)

* integrate weights manual selection bitwidth

* fix comment in test code

* Fix in set_node_quantization_config.py

* debugging version(040700gouda)

* debugging version(040701gouda)

* delete debug-print

* debugging version(040900gouda)

* revert generate_test_tpc.py

* fixed feedback comments.

* delete old tests for manual_weights_bitwidth_selection

* remove unnecessary changes.

* add typehint and description.

* fixed PR comments.

* remove unnecessary changes(new lines).

* fixed PR comments.

* fixed PR comments.

* fixed PR comments. (minimum tpc setting for e2e tests.)

---------

Co-authored-by: gouda-youichi <gouda.youichi@jp.panasonic.com>

Author: kawakami-masaki0

model_compression_toolkit/constants.py

@@ -63,6 +63,8 @@
 CORRECTED_BIAS_ATTRIBUTE = 'corrected_bias'
 ACTIVATION_N_BITS_ATTRIBUTE = 'activation_n_bits'
 SUPPORTED_INPUT_ACTIVATION_NBITS_ATTRIBUTE = 'supported_input_activation_n_bits'
+WEIGHTS = 'weights'
+ACTIVATION = 'activation'
 
 # Quantization Parameters Iterative Search Defaults:
 SYMMETRIC_TENSOR_N_ITER = 40

model_compression_toolkit/core/common/quantization/set_node_quantization_config.py

@@ -12,12 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-
-
 import copy
-from typing import List, Tuple,  Optional
+from typing import List, Tuple, Dict, Optional
 
-from mct_quantizers.common.constants import ACTIVATION_N_BITS
+from mct_quantizers.common.constants import WEIGHTS_N_BITS, ACTIVATION_N_BITS
+from model_compression_toolkit.constants import WEIGHTS, ACTIVATION
 from model_compression_toolkit.core.common import BaseNode
 from model_compression_toolkit.core.common.quantization.bit_width_config import BitWidthConfig
 from model_compression_toolkit.logger import Logger
@@ -30,14 +29,13 @@
 from model_compression_toolkit.core.common.quantization.quantization_config import QuantizationConfig, \
     QuantizationErrorMethod
 from model_compression_toolkit.core.common.quantization.quantization_params_fn_selection import \
-    get_activation_quantization_params_fn, get_weights_quantization_params_fn
-from model_compression_toolkit.core.common.quantization.quantization_fn_selection import \
-    get_weights_quantization_fn
+    get_activation_quantization_params_fn
 from model_compression_toolkit.target_platform_capabilities.schema.schema_functions import max_input_activation_n_bits
 from model_compression_toolkit.target_platform_capabilities.schema.mct_current_schema import OpQuantizationConfig, \
     QuantizationConfigOptions
 from model_compression_toolkit.target_platform_capabilities.targetplatform2framework.framework_quantization_capabilities import \
     FrameworkQuantizationCapabilities
+from model_compression_toolkit.core.common.graph.base_node import WeightAttrT
 
 
 def set_quantization_configuration_to_graph(graph: Graph,
@@ -66,16 +64,19 @@ 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_activation_bit_widths(graph)
+    nodes_to_manipulate_activation_bit_widths = {} if bit_width_config is None else bit_width_config.get_nodes_to_manipulate_activation_bit_widths(graph)
+    nodes_to_manipulate_weights_bit_widths = {} if bit_width_config is None else bit_width_config.get_nodes_to_manipulate_weights_bit_widths(graph)
 
     for n in graph.nodes:
+        manual_bit_width_override = {ACTIVATION: nodes_to_manipulate_activation_bit_widths.get(n),
+                                     WEIGHTS: nodes_to_manipulate_weights_bit_widths.get(n)}
         set_quantization_configs_to_node(node=n,
                                          graph=graph,
                                          quant_config=quant_config,
                                          fw_info=graph.fw_info,
                                          fqc=graph.fqc,
                                          mixed_precision_enable=mixed_precision_enable,
-                                         manual_bit_width_override=nodes_to_manipulate_bit_widths.get(n))
+                                         manual_bit_width_override=manual_bit_width_override)
     return graph
 
 
@@ -151,7 +152,7 @@ def set_quantization_configs_to_node(node: BaseNode,
                                      fw_info: FrameworkInfo,
                                      fqc: FrameworkQuantizationCapabilities,
                                      mixed_precision_enable: bool = False,
-                                     manual_bit_width_override: Optional[int] = None):
+                                     manual_bit_width_override: Optional[Dict] = None):
     """
     Create and set quantization configurations to a node (for both weights and activation).
 
@@ -167,8 +168,11 @@ def set_quantization_configs_to_node(node: BaseNode,
     node_qc_options = node.get_qco(fqc)
     base_config, node_qc_options_list = filter_node_qco_by_graph(node, fqc, graph, node_qc_options)
 
-    # If a manual_bit_width_override is given, filter node_qc_options_list to retain only the options with activation bits equal to manual_bit_width_override,
+    # If a manual_bit_width_override is given, filter node_qc_options_list to retain only the options with activation and weights bits equal to manual_bit_width_override,
     # and update base_config accordingly.
+    if manual_bit_width_override is None:
+        manual_bit_width_override = {ACTIVATION: None, WEIGHTS: None}
+    
     base_config, node_qc_options_list = filter_qc_options_with_manual_bit_width(
         node=node,
         node_qc_options_list=node_qc_options_list,
@@ -322,7 +326,7 @@ def filter_qc_options_with_manual_bit_width(
         node: BaseNode,
         node_qc_options_list: List[OpQuantizationConfig],
         base_config: OpQuantizationConfig,
-        manual_bit_width_override: Optional[int],
+        manual_bit_width_override: Optional[Dict],
         mixed_precision_enable: bool) -> Tuple[OpQuantizationConfig, List[OpQuantizationConfig]]:
     """
     Update the quantization configurations for a node, allowing manual bit-width overrides if specified.
@@ -331,36 +335,169 @@ def filter_qc_options_with_manual_bit_width(
         node (BaseNode): A node to set quantization configuration candidates to.
         node_qc_options_list (List[OpQuantizationConfig]): List of quantization configs for the node.
         base_config (OpQuantizationConfig): Base quantization config for the node.
-        manual_bit_width_override (Optional[int]): Specifies a custom bit-width to override the node's activation bit-width.
+        manual_bit_width_override (Optional[Dict]): Specifies a custom bit-width to override the node's activation and weights bit-width.
         mixed_precision_enable (bool): Whether mixed precision is enabled.
 
     Returns:
         Tuple[OpQuantizationConfig, List[OpQuantizationConfig]]: The updated base configuration and the filtered list of quantization configs.
     """
-    if manual_bit_width_override is None:
+    base_config, node_qc_options_list = filter_activation_qc_options_with_manual_bit_width(node,
+                                                                                           node_qc_options_list,
+                                                                                           base_config,
+                                                                                           manual_bit_width_override.get(ACTIVATION),
+                                                                                           mixed_precision_enable)
+
+    base_config, node_qc_options_list = filter_weights_qc_options_with_manual_bit_width(node,
+                                                                                        node_qc_options_list,
+                                                                                        base_config,
+                                                                                        manual_bit_width_override.get(WEIGHTS),
+                                                                                        mixed_precision_enable)
+    return base_config, node_qc_options_list
+
+
+def filter_activation_qc_options_with_manual_bit_width(
+        node: BaseNode,
+        node_qc_options_list: List[OpQuantizationConfig],
+        base_config: OpQuantizationConfig,
+        activation_manual_bit_width_override: Optional[int],
+        mixed_precision_enable: bool) -> Tuple[OpQuantizationConfig, List[OpQuantizationConfig]]:
+    """
+    Update the activation quantization configurations for a node, allowing manual bit-width overrides if specified.
+
+    Args:
+        node (BaseNode): A node to set quantization configuration candidates to.
+        node_qc_options_list (List[OpQuantizationConfig]): List of quantization configs for the node.
+        base_config (OpQuantizationConfig): Base quantization config for the node.
+        activation_manual_bit_width_override (Optional[Dict]): Specifies a custom bit-width to override the node's activation bit-width.
+        mixed_precision_enable (bool): Whether mixed precision is enabled.
+
+    Returns:
+        Tuple[OpQuantizationConfig, List[OpQuantizationConfig]]: The updated base configuration and the filtered list of quantization configs.
+    """
+    if activation_manual_bit_width_override is None:
         return base_config, node_qc_options_list
 
-    # Filter node_qc_options_list to retain only the options with activation bits equal to manual_bit_width_override.
+    # Filter node_qc_options_list to retain only the options with activation bits equal to activation_manual_bit_width_override.
     node_qc_options_list = [op_cfg for op_cfg in node_qc_options_list if
-                                manual_bit_width_override == op_cfg.activation_n_bits]
-
+                                activation_manual_bit_width_override == op_cfg.activation_n_bits]
     if len(node_qc_options_list) == 0:
-        Logger.critical(f"Manually selected activation bit-width {manual_bit_width_override} is invalid for node {node}.")
+        Logger.critical(f"Manually selected activation bit-width {activation_manual_bit_width_override} is invalid for node {node}.")
     else:
         # Update the base_config to one of the values from the filtered node_qc_options_list.
-        # First, check if a configuration similar to the original base_config but with activation bits equal to manual_bit_width_override exists.
+        # First, check if a configuration similar to the original base_config but with activation bits equal to activation_manual_bit_width_override exists.
         # If it does, use it as the base_config. If not, choose a different configuration from node_qc_options_list.
-        Logger.info(f"Setting node {node} bit-width to manually selected bit-width: {manual_bit_width_override} bits.")
-        updated_base_config = base_config.clone_and_edit({ACTIVATION_N_BITS, manual_bit_width_override})
+        Logger.info(f"Setting node {node} bit-width to manually selected bit-width: {activation_manual_bit_width_override} bits.")
+        updated_base_config = base_config.clone_and_edit({ACTIVATION_N_BITS, activation_manual_bit_width_override})
         if updated_base_config in node_qc_options_list:
-            # If a base_config with the specified manual_bit_width_override exists in the node_qc_options_list,
+            # If a base_config with the specified activation_manual_bit_width_override exists in the node_qc_options_list,
             # point the base_config to this option.
             base_config = node_qc_options_list[node_qc_options_list.index(updated_base_config)]
         else:
             # Choose a different configuration from node_qc_options_list. If multiple options exist, issue a warning.
             base_config = node_qc_options_list[0]
             if len(node_qc_options_list) > 0 and not mixed_precision_enable:
                 Logger.info(
-                    f"Request received to select {manual_bit_width_override} activation bits. However, the base configuration for layer type {node.type} is missing in the node_qc_options_list."
-                    f" Overriding base_config with an option that uses {manual_bit_width_override} bit activations.")  # pragma: no cover
-    return base_config, node_qc_options_list
+                    f"Request received to select {activation_manual_bit_width_override} activation bits. However, the base configuration for layer type {node.type} is missing in the node_qc_options_list."
+                    f" Overriding base_config with an option that uses {activation_manual_bit_width_override} bit activations.")  # pragma: no cover
+
+    return base_config, node_qc_options_list
+
+
+def filter_weights_qc_options_with_manual_bit_width(
+        node: BaseNode,
+        node_qc_options_list: List[OpQuantizationConfig],
+        base_config: OpQuantizationConfig,
+        weights_manual_bit_width_override: Optional[Tuple[int, WeightAttrT]],
+        mixed_precision_enable: bool) -> Tuple[OpQuantizationConfig, List[OpQuantizationConfig]]:
+    """
+    Update the weights quantization configurations for a node, allowing manual bit-width overrides if specified.
+
+    Args:
+        node (BaseNode): A node to set quantization configuration candidates to.
+        node_qc_options_list (List[OpQuantizationConfig]): List of quantization configs for the node.
+        base_config (OpQuantizationConfig): Base quantization config for the node.
+        weights_manual_bit_width_override (Optional[[int, WeightAttrT]]): Specifies a custom bit-width to override the node's weights bit-width.
+        mixed_precision_enable (bool): Whether mixed precision is enabled.
+
+    Returns:
+        Tuple[OpQuantizationConfig, List[OpQuantizationConfig]]: The updated base configuration and the filtered list of quantization configs.
+    """
+    if not weights_manual_bit_width_override:
+        return base_config, node_qc_options_list
+
+    # Filter node_qc_options_list to retain only the options with weights bits equal to weights_manual_bit_width_override.
+    node_qc_options_weights_list = _filter_options(node_qc_options_list, weights_manual_bit_width_override)
+
+    if len(node_qc_options_weights_list) == 0:
+        Logger.critical(f"Manually selected weights bit-width {weights_manual_bit_width_override} is invalid for node {node}.")
+    else:
+        # Update the base_config to one of the values from the filtered node_qc_options_list.
+        # First, check if a configuration similar to the original base_config but with weights bits equal to weights_manual_bit_width_override exists.
+        # If it does, use it as the base_config. If not, choose a different configuration from node_qc_options_list.
+        updated_base_config = base_config.clone_and_edit()
+
+        for bit_width, attr in weights_manual_bit_width_override:
+            Logger.info(f"Setting node {node} bit-width to manually selected {attr} bit-width: {bit_width} bits.")
+            updated_base_config = updated_base_config.clone_and_edit(attr_to_edit={attr : {WEIGHTS_N_BITS: bit_width}})
+
+        if updated_base_config in node_qc_options_weights_list:
+            # If a base_config with the specified weights_manual_bit_width_override exists in the node_qc_options_list,
+            # point the base_config to this option.
+            base_config = node_qc_options_weights_list[node_qc_options_weights_list.index(updated_base_config)]
+        else:
+            # Choose a different configuration from node_qc_options_list. If multiple options exist, issue a warning.
+            base_config = node_qc_options_weights_list[0]
+            if len(node_qc_options_weights_list) > 0 and not mixed_precision_enable:
+                Logger.info(
+                    f"Request received to select weights bit-widths {weights_manual_bit_width_override}."
+                    f"However, the base configuration for layer type {node.type} is missing in the node_qc_options_list."
+                    f" Overriding base_config with an option that uses manually selected weights bit-widths {weights_manual_bit_width_override}.")  # pragma: no cover
+
+    return base_config, node_qc_options_weights_list
+
+
+def _is_valid_option(
+        op_cfg: OpQuantizationConfig,
+        attr: WeightAttrT,
+        bit_width: int) -> bool:
+    """
+    Judge whether the specified option is valid based on the specified attribute and bit width.
+
+    Args:
+        op_cfg (OpQuantizationConfig): The quantization configuration to be judged.
+        attr (WeightAttrT): The filtered node's attributes to apply bit-width manipulation to.
+        bit_width (int): The bit width to be applied to the selected nodes.
+
+    Returns:
+        Result to judge whether the specified option is valid based on the specified attribute and bit width
+    """
+    weights_attrs = op_cfg.attr_weights_configs_mapping.keys()
+
+    if attr not in weights_attrs:
+        return False
+
+    weights_n_bits = op_cfg.attr_weights_configs_mapping[attr].weights_n_bits
+    return weights_n_bits == bit_width
+
+
+def _filter_options(
+        node_qc_options_list: List[OpQuantizationConfig],
+        weights_manual_bit_width_override: Tuple[int, WeightAttrT]) -> List[OpQuantizationConfig]:
+    """
+    Filter the options based on the specified bit width and attribute.
+
+    Args:
+        node_qc_options_list (List[OpQuantizationConfig]): List of quantization configs for the node.
+        weights_manual_bit_width_override (Tuple[int, WeightAttrT])): Specifies a custom bit-width to override the node's weights bit-width.
+
+    Returns:
+        List[OpQuantizationConfig]: Filtered the options based on the specified bit width and attribute.
+    """
+    filtered_options = []
+
+    for bit_width, attr in weights_manual_bit_width_override:
+        for op_cfg in node_qc_options_list:
+            if _is_valid_option(op_cfg, attr, bit_width):
+                filtered_options.append(op_cfg)
+
+    return filtered_options