Use float mp model baseline instead of maxbit configuration (#1441)

* float mp model for sensitivity - disable all quantizers except current candidate, instead of using maxbit config
* ensure non-negative value in similarity analyzer (due to float precision)
* add normalization method and epsilon for sensitivity metric to MP config
* print only changed by refinement solutions instead of all
* remove support for non-configurable quant layer in mp back2framework

Author: irenaby

model_compression_toolkit/core/common/mixed_precision/mixed_precision_quantization_config.py

@@ -14,11 +14,23 @@
 # ==============================================================================
 
 from dataclasses import dataclass, field
+from enum import Enum
 from typing import List, Callable, Optional
 from model_compression_toolkit.constants import MP_DEFAULT_NUM_SAMPLES, ACT_HESSIAN_DEFAULT_BATCH_SIZE
 from model_compression_toolkit.core.common.mixed_precision.distance_weighting import MpDistanceWeighting
 
 
+class MpMetricNormalization(Enum):
+    """
+    MAXBIT: normalize sensitivity metrics of layer candidates by max-bitwidth candidate (of that layer).
+    MINBIT: normalize sensitivity metrics of layer candidates by min-bitwidth candidate (of that layer).
+    NONE: no normalization.
+    """
+    MAXBIT = 'MAXBIT'
+    MINBIT = 'MINBIT'
+    NONE = 'NONE'
+
+
 @dataclass
 class MixedPrecisionQuantizationConfig:
     """
@@ -27,7 +39,6 @@ class MixedPrecisionQuantizationConfig:
     Args:
         compute_distance_fn (Callable): Function to compute a distance between two tensors. If None, using pre-defined distance methods based on the layer type for each layer.
         distance_weighting_method (MpDistanceWeighting): MpDistanceWeighting enum value that provides a function to use when weighting the distances among different layers when computing the sensitivity metric.
-        custom_metric_fn (Callable): Function to compute a custom metric. As input gets the model_mp and returns a float value for metric. If None, uses interest point metric.
         num_of_images (int): Number of images to use to evaluate the sensitivity of a mixed-precision model comparing to the float model.
         configuration_overwrite (List[int]): A list of integers that enables overwrite of mixed precision with a predefined one.
         num_interest_points_factor (float): A multiplication factor between zero and one (represents percentage) to reduce the number of interest points used to calculate the distance metric.
@@ -36,11 +47,16 @@ class MixedPrecisionQuantizationConfig:
         refine_mp_solution (bool): Whether to try to improve the final mixed-precision configuration using a greedy algorithm that searches layers to increase their bit-width, or not.
         metric_normalization_threshold (float): A threshold for checking the mixed precision distance metric values, In case of values larger than this threshold, the metric will be scaled to prevent numerical issues.
         hessian_batch_size (int): The Hessian computation batch size. used only if using mixed precision with Hessian-based objective.
-    """
+        metric_normalization (MpMetricNormalization): Metric normalization method.
+        metric_epsilon (float | None): ensure minimal distance between the metric for any non-max-bidwidth candidate
+          and a max-bitwidth candidate, i.e. metric(non-max-bitwidth) >= metric(max-bitwidth) + epsilon.
+          If none, the computed metrics are used as is.
+        custom_metric_fn (Callable): Function to compute a custom metric. As input gets the model_mp and returns a
+          float value for metric. If None, uses interest point metric.
 
+    """
     compute_distance_fn: Optional[Callable] = None
     distance_weighting_method: MpDistanceWeighting = MpDistanceWeighting.AVG
-    custom_metric_fn: Optional[Callable] = None
     num_of_images: int = MP_DEFAULT_NUM_SAMPLES
     configuration_overwrite: Optional[List[int]] = None
     num_interest_points_factor: float = field(default=1.0, metadata={"description": "Should be between 0.0 and 1.0"})
@@ -49,6 +65,9 @@ class MixedPrecisionQuantizationConfig:
     refine_mp_solution: bool = True
     metric_normalization_threshold: float = 1e10
     hessian_batch_size: int = ACT_HESSIAN_DEFAULT_BATCH_SIZE
+    metric_normalization: MpMetricNormalization = MpMetricNormalization.NONE
+    metric_epsilon: Optional[float] = 1e-6
+    custom_metric_fn: Optional[Callable] = None
     _is_mixed_precision_enabled: bool = field(init=False, default=False)
 
     def __post_init__(self):

model_compression_toolkit/core/common/mixed_precision/mixed_precision_search_facade.py

@@ -98,10 +98,11 @@ def search_bit_width(graph: Graph,
 
     # Search manager and LP are highly coupled, so LP search method was moved inside search manager.
     search_manager = MixedPrecisionSearchManager(graph,
-                                                 fw_info,
-                                                 fw_impl,
-                                                 se,
-                                                 target_resource_utilization)
+                                                 fw_info=fw_info,
+                                                 fw_impl=fw_impl,
+                                                 sensitivity_evaluator=se,
+                                                 target_resource_utilization=target_resource_utilization,
+                                                 mp_config=mp_config)
     nodes_bit_cfg = search_manager.search()
 
     graph.skip_validation_check = False

model_compression_toolkit/core/common/mixed_precision/mixed_precision_search_manager.py

@@ -12,14 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+import os
+
 import itertools
 
 import copy
 from collections import defaultdict
 
 from tqdm import tqdm
 
-from typing import Dict, List, Tuple, Optional
+from typing import Dict, List, Tuple, Optional, Set
 
 import numpy as np
 
@@ -40,6 +42,8 @@
 from model_compression_toolkit.core.common.mixed_precision.sensitivity_evaluation import SensitivityEvaluation
 from model_compression_toolkit.core.common.substitutions.apply_substitutions import substitute
 from model_compression_toolkit.logger import Logger
+from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import \
+    MixedPrecisionQuantizationConfig, MpMetricNormalization
 
 
 class MixedPrecisionSearchManager:
@@ -52,7 +56,8 @@ def __init__(self,
                  fw_info: FrameworkInfo,
                  fw_impl: FrameworkImplementation,
                  sensitivity_evaluator: SensitivityEvaluation,
-                 target_resource_utilization: ResourceUtilization):
+                 target_resource_utilization: ResourceUtilization,
+                 mp_config: MixedPrecisionQuantizationConfig):
         """
 
         Args:
@@ -74,21 +79,21 @@ def __init__(self,
 
         self.sensitivity_evaluator = sensitivity_evaluator
         self.target_resource_utilization = target_resource_utilization
+        self.mp_config = mp_config
 
         self.mp_topo_configurable_nodes = self.mp_graph.get_configurable_sorted_nodes(fw_info)
 
         self.ru_targets = target_resource_utilization.get_restricted_targets()
-        self.ru_helper = MixedPrecisionRUHelper(self.original_graph, fw_info, fw_impl)
+        self.orig_graph_ru_helper = MixedPrecisionRUHelper(self.original_graph, fw_info, fw_impl)
 
         self.min_ru_config: Dict[BaseNode, int] = self.mp_graph.get_min_candidates_config(fw_info)
-        self.max_ru_config: Dict[BaseNode, int] = self.mp_graph.get_max_candidates_config(fw_info)
 
-        self.config_reconstruction_helper = ConfigReconstructionHelper(self.original_graph)
+        self.config_reconstructor = None
+        orig_min_config = self.min_ru_config
         if self.using_virtual_graph:
-            real_min_ru_config = self.config_reconstruction_helper.reconstruct_full_configuration(self.min_ru_config)
-            self.min_ru = self.ru_helper.compute_utilization(self.ru_targets, real_min_ru_config)
-        else:
-            self.min_ru = self.ru_helper.compute_utilization(self.ru_targets, self.min_ru_config)
+            self.config_reconstructor = ConfigReconstructionHelper(self.original_graph)
+            orig_min_config = self.config_reconstructor.reconstruct_full_configuration(self.min_ru_config)
+        self.min_ru = self.orig_graph_ru_helper.compute_utilization(self.ru_targets, orig_min_config)
 
     def search(self) -> Dict[BaseNode, int]:
         """
@@ -100,7 +105,7 @@ def search(self) -> Dict[BaseNode, int]:
         mp_config = self._prepare_and_run_solver()
 
         if self.using_virtual_graph:
-            mp_config = self.config_reconstruction_helper.reconstruct_full_configuration(mp_config)
+            mp_config = self.config_reconstructor.reconstruct_full_configuration(mp_config)
 
         return mp_config
 
@@ -143,61 +148,64 @@ def _get_relative_ru_constraint_per_mem_element(self) -> Dict[RUTarget, np.ndarr
                              f"following targets: {unsatisfiable_targets}")
         return rel_target_ru
 
-    def _build_sensitivity_mapping(self, eps: float = 1e-6) -> Dict[BaseNode, List[float]]:
+    def _build_sensitivity_mapping(self) -> Dict[BaseNode, List[float]]:
         """
         This function measures the sensitivity of a change in a bitwidth of a layer on the entire model.
 
-        Args:
-            eps: if sensitivity for a non-max candidate is lower than for a max candidate, we set it to
-              sensitivity of a max candidate + epsilon.
-
         Returns:
             Mapping from nodes to their bitwidth candidates sensitivity.
         """
-
         Logger.info('Starting to evaluate metrics')
-
-        orig_sorted_nodes = self.original_graph.get_configurable_sorted_nodes(self.fw_info)
-
-        def topo_cfg(cfg: dict) -> list:
-            topo_cfg = [cfg[n] for n in orig_sorted_nodes]
-            assert len(topo_cfg) == len(cfg)
-            return topo_cfg
-
-        def compute_metric(cfg, node_idx=None, baseline_cfg=None):
-            return self.sensitivity_evaluator.compute_metric(topo_cfg(cfg),
-                                                             node_idx,
-                                                             topo_cfg(baseline_cfg) if baseline_cfg else None)
-
-        if self.using_virtual_graph:
-            origin_max_config = self.config_reconstruction_helper.reconstruct_full_configuration(self.max_ru_config)
-            max_config_value = compute_metric(origin_max_config)
-        else:
-            max_config_value = compute_metric(self.max_ru_config)
+        norm_method = self.mp_config.metric_normalization
+        eps = self.mp_config.metric_epsilon
+
+        verbose = 'VERBOSE_MP_METRIC' in os.environ
+
+        def normalize(node_candidates_metrics, max_ind):
+            if norm_method == MpMetricNormalization.NONE:
+                return node_candidates_metrics
+            if norm_method == MpMetricNormalization.MAXBIT:
+                ref_ind = max_ind
+            elif norm_method == MpMetricNormalization.MINBIT:
+                ref_ind = node.find_min_candidate_index()
+            else:  # pragma: no cover
+                raise ValueError(f'Unexpected MpMetricNormalization mode {norm_method}')
+            normalized_metrics = node_candidates_metrics / node_candidates_metrics[ref_ind]
+            if verbose and not np.array_equal(normalized_metrics, node_candidates_metrics):
+                print(f'{"normalized metric:":25}', candidates_sensitivity)
+            return normalized_metrics
+
+        def ensure_maxbit_minimal_metric(node_candidates_metrics, max_ind):
+            if eps is None:
+                return node_candidates_metrics
+            # We want maxbit configuration to have the minimal distance metric (so that optimization objective
+            # doesn't prefer lower bits). If we got a smaller metric for non-maxbit, we update it to metric(maxbit)+eps.
+            max_val = node_candidates_metrics[max_ind]
+            metrics = np.maximum(node_candidates_metrics, max_val + eps)
+            metrics[max_ind] = max_val
+            if verbose and not np.array_equal(metrics, node_candidates_metrics):
+                print(f'{"eps-adjusted metric:":25}', candidates_sensitivity)
+            return metrics
 
         layer_to_metrics_mapping = defaultdict(list)
         for node_idx, node in tqdm(enumerate(self.mp_topo_configurable_nodes)):
+            candidates_sensitivity = np.empty(len(node.candidates_quantization_cfg))
             for bitwidth_idx, _ in enumerate(node.candidates_quantization_cfg):
-                if self.max_ru_config[node] == bitwidth_idx:
-                    # This is a computation of the metric for the max configuration, assign pre-calculated value
-                    layer_to_metrics_mapping[node].append(max_config_value)
-                    continue
-
-                # Create a configuration that differs at one layer only from the baseline model
-                mp_model_configuration = self.max_ru_config.copy()
-                mp_model_configuration[node] = bitwidth_idx
-
-                # Build a distance matrix using the function we got from the framework implementation.
                 if self.using_virtual_graph:
-                    # Reconstructing original graph's configuration from virtual graph's configuration
-                    orig_mp_config = self.config_reconstruction_helper.reconstruct_full_configuration(mp_model_configuration)
-                    changed_nodes = [orig_sorted_nodes.index(n) for n, ind in orig_mp_config.items()
-                                     if origin_max_config[n] != ind]
-                    metric_value = compute_metric(orig_mp_config, changed_nodes, origin_max_config)
+                    a_cfg, w_cfg = self.config_reconstructor.reconstruct_separate_aw_configs({node: bitwidth_idx})
                 else:
-                    metric_value = compute_metric(mp_model_configuration, [node_idx], self.max_ru_config)
-                metric_value = max(metric_value, max_config_value + eps)
-                layer_to_metrics_mapping[node].append(metric_value)
+                    a_cfg = {node: bitwidth_idx} if node.has_configurable_activation() else {}
+                    w_cfg = {node: bitwidth_idx} if node.has_any_configurable_weight() else {}
+                candidates_sensitivity[bitwidth_idx] = self.sensitivity_evaluator.compute_metric(
+                    mp_a_cfg={n.name: ind for n, ind in a_cfg.items()},
+                    mp_w_cfg={n.name: ind for n, ind in w_cfg.items()}
+                )
+            if verbose:
+                print(f'{node.name}\n{"raw metric:":25}', candidates_sensitivity)
+            max_ind = node.find_max_candidate_index()
+            candidates_sensitivity = normalize(candidates_sensitivity, max_ind)
+            candidates_sensitivity = ensure_maxbit_minimal_metric(candidates_sensitivity, max_ind)
+            layer_to_metrics_mapping[node] = candidates_sensitivity
 
         # Finalize distance metric mapping
         self._finalize_distance_metric(layer_to_metrics_mapping)
@@ -244,8 +252,9 @@ def _compute_relative_ru_matrices(self) -> Dict[RUTarget, np.ndarray]:
                 else:
                     cfg = self.min_ru_config.copy()
                     cfg[node] = candidate_idx
-                    real_cfg = self.config_reconstruction_helper.reconstruct_full_configuration(cfg)
-                    candidate_rus = self.ru_helper.compute_utilization(self.ru_targets, real_cfg)
+                    if self.using_virtual_graph:
+                        cfg = self.config_reconstructor.reconstruct_full_configuration(cfg)
+                    candidate_rus = self.orig_graph_ru_helper.compute_utilization(self.ru_targets, cfg)
 
                 for target, ru in candidate_rus.items():
                     rus_per_candidate[target].append(ru)
@@ -283,8 +292,8 @@ def compute_resource_utilization_for_config(self, config: Dict[BaseNode, int]) -
         with the given config.
 
         """
-        act_qcs, w_qcs = self.ru_helper.get_quantization_candidates(config)
-        ru = self.ru_helper.ru_calculator.compute_resource_utilization(
+        act_qcs, w_qcs = self.orig_graph_ru_helper.get_quantization_candidates(config)
+        ru = self.orig_graph_ru_helper.ru_calculator.compute_resource_utilization(
             target_criterion=TargetInclusionCriterion.AnyQuantized, bitwidth_mode=BitwidthMode.QCustom, act_qcs=act_qcs,
             w_qcs=w_qcs, ru_targets=self.ru_targets, allow_unused_qcs=True)
         return ru
@@ -303,7 +312,7 @@ def _finalize_distance_metric(self, layer_to_metrics_mapping: Dict[BaseNode, Lis
         # normalize metric for numerical stability
         max_dist = max(itertools.chain.from_iterable(layer_to_metrics_mapping.values()))
 
-        if max_dist >= self.sensitivity_evaluator.quant_config.metric_normalization_threshold:
+        if max_dist >= self.mp_config.metric_normalization_threshold:
             Logger.warning(f"The mixed precision distance metric values indicate a large error in the quantized model."
                            f"this can cause numerical issues."
                            f"The program will proceed with mixed precision search after scaling the metric values,"
@@ -387,7 +396,9 @@ def reconstruct_full_configuration(self,
 
         return orig_cfg
 
-    def reconstruct_separate_aw_configs(self, virtual_cfg: Dict[BaseNode, int], include_non_configurable: bool) \
+    def reconstruct_separate_aw_configs(self,
+                                        virtual_cfg: Dict[BaseNode, int],
+                                        include_non_configurable: bool = False) \
             -> Tuple[Dict[BaseNode, int], Dict[BaseNode, int]]:
         """
         Retrieves original activation and weights nodes and corresponding candidates for a given configuration of the