add support for custom metric function for mixed precision

Author: itai-berman

model_compression_toolkit/core/common/mixed_precision/mixed_precision_quantization_config.py

@@ -39,6 +39,7 @@ class MixedPrecisionQuantizationConfig:
 
     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"})

model_compression_toolkit/core/common/mixed_precision/mixed_precision_search_manager.py

@@ -162,9 +162,15 @@ def topo_cfg(cfg: dict) -> list:
             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.sensitivity_evaluator.quant_config.custom_metric_fn is None:
+                return self.sensitivity_evaluator.compute_metric(topo_cfg(cfg),
+                                                                 node_idx,
+                                                                 topo_cfg(baseline_cfg) if baseline_cfg else None)
+            else:
+                return self.sensitivity_evaluator.compute_custom_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_config_from_virtual_graph(
                 self.max_ru_config)

model_compression_toolkit/core/common/mixed_precision/sensitivity_evaluation.py

@@ -89,13 +89,19 @@ def __init__(self,
         self.interest_points = get_mp_interest_points(graph,
                                                       fw_impl.count_node_for_mixed_precision_interest_points,
                                                       quant_config.num_interest_points_factor)
+        # If using a custom metric - return only model outputs
+        if self.quant_config.custom_metric_fn is not None:
+            self.interest_points = []
 
         # We use normalized MSE when not running hessian-based. For Hessian-based normalized MSE is not needed
         # because hessian weights already do normalization.
         use_normalized_mse = self.quant_config.use_hessian_based_scores is False
         self.ips_distance_fns, self.ips_axis = self._init_metric_points_lists(self.interest_points, use_normalized_mse)
 
         self.output_points = get_output_nodes_for_metric(graph)
+        # If using a custom metric - return all model outputs
+        if self.quant_config.custom_metric_fn is not None:
+            self.output_points = [n.node for n in graph.get_outputs()]
         self.out_ps_distance_fns, self.out_ps_axis = self._init_metric_points_lists(self.output_points,
                                                                                     use_normalized_mse)
 
@@ -187,6 +193,38 @@ def compute_metric(self,
         return self._compute_mp_distance_measure(ipts_distances, out_pts_distances,
                                                  self.quant_config.distance_weighting_method)
 
+    def compute_custom_metric(self,
+                              mp_model_configuration: List[int],
+                              node_idx: List[int] = None,
+                              baseline_mp_configuration: List[int] = None) -> float:
+        """
+        Compute the sensitivity metric of the MP model for a given configuration (the sensitivity
+        is computed on a custom function).
+
+        Args:
+            mp_model_configuration: Bitwidth configuration to use to configure the MP model.
+            node_idx: A list of nodes' indices to configure (instead of using the entire mp_model_configuration).
+            baseline_mp_configuration: A mixed-precision configuration to set the model back to after modifying it to
+                compute the metric for the given configuration.
+
+        Returns:
+            The sensitivity metric of the MP model for a given configuration.
+        """
+
+        # Configure MP model with the given configuration.
+        self._configure_bitwidths_model(mp_model_configuration,
+                                        node_idx)
+
+        # Compute the distance metric
+        sensitivity_metric = self.quant_config.custom_metric_fn(self.model_mp)
+
+        # Configure MP model back to the same configuration as the baseline model if baseline provided
+        if baseline_mp_configuration is not None:
+            self._configure_bitwidths_model(baseline_mp_configuration,
+                                            node_idx)
+
+        return sensitivity_metric
+
     def _init_baseline_tensors_list(self):
         """
         Evaluates the baseline model on all images and returns the obtained lists of tensors in a list for later use.