mct-nightly 2.3.0.20250602.610__py3-none-any.whl → 2.3.0.20250604.611__py3-none-any.whl

{mct_nightly-2.3.0.20250602.610.dist-info → mct_nightly-2.3.0.20250604.611.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mct-nightly
-Version: 2.3.0.20250602.610
+Version: 2.3.0.20250604.611
 Summary: A Model Compression Toolkit for neural networks
 Author-email: ssi-dnn-dev@sony.com
 Classifier: Programming Language :: Python :: 3
@@ -236,6 +236,6 @@ MCT is licensed under Apache License Version 2.0. By contributing to the project
 
 [3] [TORCHVISION.MODELS](https://pytorch.org/vision/stable/models.html) 
 
-[4] Gordon, O., Cohen, E., Habi, H. V., & Netzer, A., 2024. [EPTQ: Enhanced Post-Training Quantization via Hessian-guided Network-wise Optimization, European Conference on Computer Vision Workshop 2024, Computational Aspects of Deep Learning (CADL)](https://arxiv.org/abs/2309.11531)
+[4] Gordon, O., Cohen, E., Habi, H.V., Netzer, A. (2025). [EPTQ: Enhanced Post-Training Quantization via Hessian-guided Network-wise Optimization – ECCV 2024 Workshops](https://link.springer.com/chapter/10.1007/978-3-031-91979-4_13)
 
 [5] Dikstein, L., Lapid, A., Netzer, A., & Habi, H. V., 2024. [Data Generation for Hardware-Friendly Post-Training Quantization, Accepted to IEEE/CVF Winter Conference on Applications of Computer Vision (WACV) 2025](https://openaccess.thecvf.com/content/WACV2025/papers/Dikstein_Data_Generation_for_Hardware-Friendly_Post-Training_Quantization_WACV_2025_paper.pdf)

{mct_nightly-2.3.0.20250602.610.dist-info → mct_nightly-2.3.0.20250604.611.dist-info}/RECORD

@@ -1,11 +1,11 @@
-mct_nightly-2.3.0.20250602.610.dist-info/licenses/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
-model_compression_toolkit/__init__.py,sha256=UJNGJtOutpsyDOmEKhw0oA_frEZl4iGhFMqRsWW16Ac,1557
+mct_nightly-2.3.0.20250604.611.dist-info/licenses/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
+model_compression_toolkit/__init__.py,sha256=MNjhJRHWg20p0VjHTzI6AhnDetVQ7Oy8YQrAOLRA-4Y,1557
 model_compression_toolkit/constants.py,sha256=KNgiNLpsMgSYyXMNEbHXd4bFNerQc1D6HH3vpbUq_Gs,4086
 model_compression_toolkit/defaultdict.py,sha256=LSc-sbZYXENMCw3U9F4GiXuv67IKpdn0Qm7Fr11jy-4,2277
 model_compression_toolkit/logger.py,sha256=L3q7tn3Uht0i_7phnlOWMR2Te2zvzrt2HOz9vYEInts,4529
 model_compression_toolkit/metadata.py,sha256=x_Bk4VpzILdsFax6--CZ3X18qUTP28sbF_AhoQW8dNc,4003
 model_compression_toolkit/verify_packages.py,sha256=l0neIRr8q_QwxmuiTI4vyCMDISDedK0EihjEQUe66tE,1319
-model_compression_toolkit/core/__init__.py,sha256=phfdtc09uruSyOpWRaUMUeMNRSwYB5q9NBus3cqcjIM,2113
+model_compression_toolkit/core/__init__.py,sha256=HNverPpoqEyFKTa7iEdOqqY2P0Gq-7GMejNOi6ZPcQs,2042
 model_compression_toolkit/core/analyzer.py,sha256=X-2ZpkH1xdXnISnw1yJvXnvV-ssoUh-9LkLISSWNqiY,3691
 model_compression_toolkit/core/graph_prep_runner.py,sha256=C6eUTd-fcgxk0LUbt51gFZwmyDDDEB8-9Q4kr9ujYvI,11555
 model_compression_toolkit/core/quantization_prep_runner.py,sha256=DPevqQ8brkdut8K5f5v9g5lbT3r1GSmhLAk3NkL40Fg,6593
@@ -60,15 +60,15 @@ model_compression_toolkit/core/common/matchers/edge_matcher.py,sha256=bS9KIBhB6Y
 model_compression_toolkit/core/common/matchers/function.py,sha256=kMwcinxn_PInvetNh_L_lqGXT1hoi3f97PqBpjqfXoA,1773
 model_compression_toolkit/core/common/matchers/node_matcher.py,sha256=63cMwa5YbQ5LKZy8-KFmdchVc3N7mpDJ6fNDt_uAQsk,2745
 model_compression_toolkit/core/common/matchers/walk_matcher.py,sha256=xqfLKk6xZt72hSnND_HoX5ESOooNMypb5VOZkVsJ_nw,1111
-model_compression_toolkit/core/common/mixed_precision/__init__.py,sha256=Jm6pls3QUCMQ9d86KOYxOq05br_k130ByGHLCojIZ_M,766
+model_compression_toolkit/core/common/mixed_precision/__init__.py,sha256=Vlpo9M_1u6LHdEjYE3-wGc1esoH2NVhRzi3n_HTYvHs,789
 model_compression_toolkit/core/common/mixed_precision/bit_width_setter.py,sha256=npqLPyk5xXR11M_zdImtSALc5vJv9N4fEapaludKLBw,7139
 model_compression_toolkit/core/common/mixed_precision/configurable_quant_id.py,sha256=LLDguK7afsbN742ucLpmJr5TUfTyFpK1vbf2bpVr1v0,882
 model_compression_toolkit/core/common/mixed_precision/configurable_quantizer_utils.py,sha256=7dKMi5S0zQZ16m8NWn1XIuoXsKuZUg64G4-uK8-j1PQ,5177
 model_compression_toolkit/core/common/mixed_precision/mixed_precision_candidates_filter.py,sha256=6pLUEEIqRTVIlCYQC4JIvY55KAvuBHEX8uTOQ-1Ac4Q,3859
-model_compression_toolkit/core/common/mixed_precision/mixed_precision_quantization_config.py,sha256=rdtxPmRhjrC160O3fqAjDzGxpMeM49hYhmlnf_Kwqds,5416
+model_compression_toolkit/core/common/mixed_precision/mixed_precision_quantization_config.py,sha256=3vuhlpGvUdeGpEMfNIYONMN9NQ_VFIJykC6Gwlo-gaI,6728
 model_compression_toolkit/core/common/mixed_precision/mixed_precision_ru_helper.py,sha256=axgAypzsiCOw04ZOtOEjK4riuNsaEU2qU6KkWnEXtMo,4951
 model_compression_toolkit/core/common/mixed_precision/mixed_precision_search_facade.py,sha256=1877xOUdgpWrXWyhdX1pJOePuopq43L71WqBFMqzyR4,6418
-model_compression_toolkit/core/common/mixed_precision/mixed_precision_search_manager.py,sha256=TAxA9BKxINwUQfJpmf2Qghz-5DTbesuf1Pe1L0Tc-j4,28157
+model_compression_toolkit/core/common/mixed_precision/mixed_precision_search_manager.py,sha256=3D_gHdcxsuINiTZqEAJXbxwYDg5qqXD51k4_smmWI9M,28553
 model_compression_toolkit/core/common/mixed_precision/solution_refinement_procedure.py,sha256=MY8df-c_kITEr_7hOctaxhdiq29hSTA0La9Qo0oTJJY,9678
 model_compression_toolkit/core/common/mixed_precision/resource_utilization_tools/__init__.py,sha256=Rf1RcYmelmdZmBV5qOKvKWF575ofc06JFQSq83Jz99A,696
 model_compression_toolkit/core/common/mixed_precision/resource_utilization_tools/resource_utilization.py,sha256=PKkhc5q8pEPnNLXwo3U56EOCfYnPXIvPs0LlCGZOoKU,4426
@@ -77,8 +77,7 @@ model_compression_toolkit/core/common/mixed_precision/resource_utilization_tools
 model_compression_toolkit/core/common/mixed_precision/search_methods/__init__.py,sha256=sw7LOPN1bM82o3SkMaklyH0jw-TLGK0-fl2Wq73rffI,697
 model_compression_toolkit/core/common/mixed_precision/search_methods/linear_programming.py,sha256=6Z6nQL9UH7B8dbcUR0cuCTEYFOKZAlvOb-SCk_cAZFA,6670
 model_compression_toolkit/core/common/mixed_precision/sensitivity_eval/__init__.py,sha256=5yxITHNJcCfeGKdIpAYbNbKDoXUSvENuRQm3OQu8Qf4,697
-model_compression_toolkit/core/common/mixed_precision/sensitivity_eval/distance_weighting.py,sha256=-x8edUyudu1EAEM66AuXPtgayLpzbxoLNubfEbFM5kU,2867
-model_compression_toolkit/core/common/mixed_precision/sensitivity_eval/metric_calculators.py,sha256=W4CySFtN874npcM9j9wu1PVrv7IZHLyKdLOPrTsCNQg,22209
+model_compression_toolkit/core/common/mixed_precision/sensitivity_eval/metric_calculators.py,sha256=WhSsQDsmunfYlAU34-YUYnQXEiPBFf2Gl4FnyCn-3LE,21966
 model_compression_toolkit/core/common/mixed_precision/sensitivity_eval/sensitivity_evaluation.py,sha256=5l0qP0mZ061xh3rjqTJZcLD2mMKC-hfSnNAN0OmSusk,8938
 model_compression_toolkit/core/common/mixed_precision/sensitivity_eval/set_layer_to_bitwidth.py,sha256=Zn6SgzGLWWKmuYGHd1YtKxZdYnQWRDeXEkKlBiTbHcs,2929
 model_compression_toolkit/core/common/network_editors/__init__.py,sha256=vZmu55bYqiaOQs3AjfwWDXHmuKZcLHt-wm7uR5fPEqg,1307
@@ -135,7 +134,7 @@ model_compression_toolkit/core/common/statistics_correction/apply_activation_bia
 model_compression_toolkit/core/common/statistics_correction/apply_bias_correction_to_graph.py,sha256=b5clhUWGoDaQLn2pDCeYkV0FomVebcKS8pMXtQTTzIg,4679
 model_compression_toolkit/core/common/statistics_correction/apply_second_moment_correction_to_graph.py,sha256=C_nwhhitTd1pCto0nHZPn3fjIMOeDD7VIciumTR3s6k,5641
 model_compression_toolkit/core/common/statistics_correction/compute_activation_bias_correction_of_graph.py,sha256=zIkhOPF6K5aIgMExpD7HFT9UZSDpvXh51F6V-qZ7H-4,9048
-model_compression_toolkit/core/common/statistics_correction/compute_bias_correction_of_graph.py,sha256=LaGhYES7HgIDf9Bi2KAG_mBzAWuum0J6AGmAFPC8wwo,10478
+model_compression_toolkit/core/common/statistics_correction/compute_bias_correction_of_graph.py,sha256=Jditohcz1mGEYZTqNlzNyoRMsjZx0hsz1gkN1xux_HI,10596
 model_compression_toolkit/core/common/statistics_correction/statistics_correction.py,sha256=E0ZA4edimJwpHh9twI5gafcoJ9fX5F1JX2QUOkUOKEw,6250
 model_compression_toolkit/core/common/substitutions/__init__.py,sha256=sw7LOPN1bM82o3SkMaklyH0jw-TLGK0-fl2Wq73rffI,697
 model_compression_toolkit/core/common/substitutions/apply_substitutions.py,sha256=k-bifmakHIYZeZS-4T1QpZ1Et6AwAijMRgAKs7hmMKc,1390
@@ -530,7 +529,7 @@ model_compression_toolkit/xquant/pytorch/model_analyzer.py,sha256=b93o800yVB3Z-i
 model_compression_toolkit/xquant/pytorch/pytorch_report_utils.py,sha256=UVN_S9ULHBEldBpShCOt8-soT8YTQ5oE362y96qF_FA,3950
 model_compression_toolkit/xquant/pytorch/similarity_functions.py,sha256=CERxq5K8rqaiE-DlwhZBTUd9x69dtYJlkHOPLB54vm8,2354
 model_compression_toolkit/xquant/pytorch/tensorboard_utils.py,sha256=mkoEktLFFHtEKzzFRn_jCnxjhJolK12TZ5AQeDHzUO8,9767
-mct_nightly-2.3.0.20250602.610.dist-info/METADATA,sha256=Snekj6AEV65NneMEyhZizOHy6rmAYUscyrhiFRnkJMg,25136
-mct_nightly-2.3.0.20250602.610.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-mct_nightly-2.3.0.20250602.610.dist-info/top_level.txt,sha256=gsYA8juk0Z-ZmQRKULkb3JLGdOdz8jW_cMRjisn9ga4,26
-mct_nightly-2.3.0.20250602.610.dist-info/RECORD,,
+mct_nightly-2.3.0.20250604.611.dist-info/METADATA,sha256=vtgFVscrMOf1owGEtgvDKYIE2oCQ2ewOg0-JbeDPI0k,25087
+mct_nightly-2.3.0.20250604.611.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+mct_nightly-2.3.0.20250604.611.dist-info/top_level.txt,sha256=gsYA8juk0Z-ZmQRKULkb3JLGdOdz8jW_cMRjisn9ga4,26
+mct_nightly-2.3.0.20250604.611.dist-info/RECORD,,

model_compression_toolkit/__init__.py

@@ -27,4 +27,4 @@ from model_compression_toolkit import data_generation
 from model_compression_toolkit import pruning
 from model_compression_toolkit.trainable_infrastructure.keras.load_model import keras_load_quantized_model
 
-__version__ = "2.3.0.20250602.000610"
+__version__ = "2.3.0.20250604.000611"

model_compression_toolkit/core/__init__.py

@@ -22,8 +22,8 @@ from model_compression_toolkit.core.common.quantization.quantization_config impo
 from model_compression_toolkit.core.common.quantization.bit_width_config import BitWidthConfig
 from model_compression_toolkit.core.common.quantization.core_config import CoreConfig
 from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization import ResourceUtilization
-from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import MixedPrecisionQuantizationConfig
+from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import (
+    MixedPrecisionQuantizationConfig, MpDistanceWeighting, MpMetricNormalization)
 from model_compression_toolkit.core.keras.resource_utilization_data_facade import keras_resource_utilization_data
 from model_compression_toolkit.core.pytorch.resource_utilization_data_facade import pytorch_resource_utilization_data
-from model_compression_toolkit.core.common.mixed_precision.sensitivity_eval.distance_weighting import MpDistanceWeighting
 

model_compression_toolkit/core/common/mixed_precision/__init__.py

@@ -12,4 +12,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from .sensitivity_eval.distance_weighting import MpDistanceWeighting
+from .mixed_precision_quantization_config import MpDistanceWeighting, MpMetricNormalization

model_compression_toolkit/core/common/mixed_precision/mixed_precision_quantization_config.py

@@ -14,10 +14,25 @@
 # ==============================================================================
 
 from dataclasses import dataclass, field
-from enum import Enum
+from enum import Enum, auto
 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.sensitivity_eval.distance_weighting import MpDistanceWeighting
+
+
+class MpDistanceWeighting(Enum):
+    """
+    Defines interest points distances weighting methods.
+
+    AVG - take the average distance over all interest points.
+    LAST_LAYER - take only the distance of the last interest point.
+    EXP - weighted average with weights based on exponent of negative distances between activations of the quantized and the float models.
+    HESSIAN - weighted average with Hessians as weights.
+    """
+    AVG = auto()
+    LAST_LAYER = auto()
+    EXP = auto()
+    HESSIAN = auto()
 
 
 class MpMetricNormalization(Enum):
@@ -38,11 +53,12 @@ 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.
+        distance_weighting_method (MpDistanceWeighting): distance weighting method to use. By default, MpDistanceWeighting.AVG.
         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.
-        use_hessian_based_scores (bool): Whether to use Hessian-based scores for weighted average distance metric computation.
+        use_hessian_based_scores (bool): Whether to use Hessian-based scores for weighted average distance metric
+          computation. This is identical to passing distance_weighting_method=MpDistanceWeighting.HESSIAN.
         norm_scores (bool): Whether to normalize the returned scores for the weighted distance metric (to get values between 0 and 1).
         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.
@@ -51,12 +67,14 @@ class MixedPrecisionQuantizationConfig:
         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.
+        exp_distance_weighting_sigma (float): sigma for exponential weighting method. A distance for each interest point
+          is normalized by sigma prior to applying exponent.
         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
+    distance_weighting_method: MpDistanceWeighting = 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"})
@@ -67,6 +85,7 @@ class MixedPrecisionQuantizationConfig:
     hessian_batch_size: int = ACT_HESSIAN_DEFAULT_BATCH_SIZE
     metric_normalization: MpMetricNormalization = MpMetricNormalization.NONE
     metric_epsilon: Optional[float] = 1e-6
+    exp_distance_weighting_sigma: float = 0.1
     custom_metric_fn: Optional[Callable] = None
     _is_mixed_precision_enabled: bool = field(init=False, default=False)
 
@@ -77,6 +96,14 @@ class MixedPrecisionQuantizationConfig:
             "the base set of interest points that are required to be " \
             "used for mixed-precision metric evaluation, " \
             "thus, it should be between 0 to 1"
+        if self.use_hessian_based_scores:
+            assert self.distance_weighting_method in [None, MpDistanceWeighting.HESSIAN], \
+                f'Distance method {self.distance_weighting_method} is incompatible with use_hessian_based_scores=True'
+            self.distance_weighting_method = MpDistanceWeighting.HESSIAN
+        elif self.distance_weighting_method is None and self.custom_metric_fn is None:
+            self.distance_weighting_method = MpDistanceWeighting.AVG
+        assert self.exp_distance_weighting_sigma > 0, (f'exp_distance_weighting_sigma should be positive, but got '
+                                                       f'{self.exp_distance_weighting_sigma}')
 
     def set_mixed_precision_enable(self):
         """

model_compression_toolkit/core/common/mixed_precision/mixed_precision_search_manager.py

@@ -171,8 +171,6 @@ class MixedPrecisionSearchManager:
             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):
@@ -183,30 +181,41 @@ class MixedPrecisionSearchManager:
             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)
+        layer_to_metrics_mapping = {}
+        debug_mapping = {}
         for node_idx, node in tqdm(enumerate(self.mp_topo_configurable_nodes)):
-            candidates_sensitivity = np.empty(len(node.candidates_quantization_cfg))
+            raw_candidates_sensitivity = np.empty(len(node.candidates_quantization_cfg))
             for bitwidth_idx, _ in enumerate(node.candidates_quantization_cfg):
                 if self.using_virtual_graph:
                     a_cfg, w_cfg = self.config_reconstructor.reconstruct_separate_aw_configs({node: bitwidth_idx})
                 else:
                     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(
+                raw_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)
+            normalized_sensitivity = normalize(raw_candidates_sensitivity, max_ind)
+            candidates_sensitivity = ensure_maxbit_minimal_metric(normalized_sensitivity, max_ind)
             layer_to_metrics_mapping[node] = candidates_sensitivity
 
+            if verbose:    # pragma: no cover
+                debug_mapping[node] = {'': candidates_sensitivity}
+                if np.any(raw_candidates_sensitivity != candidates_sensitivity):
+                    debug_mapping[node]['normalized'] = normalized_sensitivity
+                    debug_mapping[node]['raw       '] = raw_candidates_sensitivity
+
+        if verbose:    # pragma: no cover
+            np.set_printoptions(precision=8, floatmode='maxprec')
+            name_len = max(len(n.name) for n in debug_mapping)
+            s = '\nMETRIC BEGIN'
+            for n, d in debug_mapping.items():
+                s += (f'\n{n.name:{name_len}}' + f'\n{" ":{name_len-10}}'.join([f'{k} {v}' for k, v in d.items()]))
+            s += '\nMETRIC END'
+            Logger.info(s)
         # Finalize distance metric mapping
         self._finalize_distance_metric(layer_to_metrics_mapping)
 
@@ -372,7 +381,7 @@ class ConfigReconstructionHelper:
             w_candidates = orig_nodes_w_candidates[orig_node]
             # find the common candidate
             common_candidates = set(a_candidates).intersection(set(w_candidates))
-            if len(common_candidates) != 1:
+            if len(common_candidates) != 1:    # pragma: no cover
                 raise ValueError(f'Expected to find exactly one candidate with the required activation and weights '
                                  f'quantization configuration for node {orig_node}. Found {len(common_candidates)}')
             # in theory it's possible that original non-configurable node gets split and each part is combined

model_compression_toolkit/core/common/mixed_precision/sensitivity_eval/metric_calculators.py

@@ -15,7 +15,7 @@
 import numpy as np
 from typing import runtime_checkable, Protocol, Callable, Any, List, Tuple
 
-from model_compression_toolkit.core import MixedPrecisionQuantizationConfig, FrameworkInfo
+from model_compression_toolkit.core import FrameworkInfo, MixedPrecisionQuantizationConfig, MpDistanceWeighting
 from model_compression_toolkit.core.common import Graph, BaseNode
 from model_compression_toolkit.core.common.hessian import HessianInfoService, HessianScoresRequest, HessianMode, \
     HessianScoresGranularity
@@ -32,7 +32,7 @@ class MetricCalculator(Protocol):
 
     def compute(self, mp_model) -> float:
         """ Compute the metric for the given model. """
-        ...
+        raise NotImplementedError    # pragma: no cover
 
 
 class CustomMetricCalculator(MetricCalculator):
@@ -66,16 +66,6 @@ class DistanceMetricCalculator(MetricCalculator):
                  fw_impl: Any,
                  hessian_info_service: HessianInfoService = None):
         """
-        Initiates all relevant objects to manage a sensitivity evaluation for MP search.
-        Create an object that allows to compute the sensitivity metric of an MP model (the sensitivity
-        is computed based on the similarity of the interest points' outputs between the MP model
-        and the float model).
-        First, we initiate a SensitivityEvaluationManager that handles the components which are necessary for
-        evaluating the sensitivity. It initializes an MP model (a model where layers that can be configured in
-        different bit-widths) and a baseline model (a float model).
-        Then, and based on the outputs of these two models (for some batches from the representative_data_gen),
-        we build a function to measure the sensitivity of a change in a bit-width of a model's layer.
-
         Args:
             graph: Graph to search for its MP configuration.
             mp_config: MP Quantization configuration for how the graph should be quantized.
@@ -91,11 +81,9 @@ class DistanceMetricCalculator(MetricCalculator):
         self.fw_info = fw_info
         self.fw_impl = fw_impl
 
-        if self.mp_config.use_hessian_based_scores:
-            if not isinstance(hessian_info_service, HessianInfoService):  # pragma: no cover
-                Logger.critical(
-                    f"When using Hessian-based approximations for sensitivity evaluation, a valid HessianInfoService object is required; found {type(hessian_info_service)}.")
-            self.hessian_info_service = hessian_info_service
+        if self.mp_config.distance_weighting_method == MpDistanceWeighting.HESSIAN:
+            assert hessian_info_service is not None, ('Expected HessianInfoService object to be passed with Hessian '
+                                                      'distance weighting')
 
         self.sorted_configurable_nodes_names = graph.get_configurable_sorted_nodes_names(self.fw_info)
 
@@ -104,17 +92,12 @@ class DistanceMetricCalculator(MetricCalculator):
         self.interest_points = self.get_mp_interest_points(graph,
                                                            fw_impl.count_node_for_mixed_precision_interest_points,
                                                            mp_config.num_interest_points_factor)
-
-        # 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.mp_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)
-
+        # exponential weighing assumes normalized distances, we only store it to be able to catch any changes
+        self.use_normalized_mse = True
+        self.ips_distance_fns, self.ips_axis = self._init_metric_points_lists(self.interest_points)
         output_points = self.get_output_nodes_for_metric(graph)
         self.all_interest_points = self.interest_points + output_points
-        self.out_ps_distance_fns, self.out_ps_axis = self._init_metric_points_lists(output_points,
-                                                                                    use_normalized_mse)
+        self.out_ps_distance_fns, self.out_ps_axis = self._init_metric_points_lists(output_points)
 
         self.ref_model, _ = fw_impl.model_builder(graph, mode=ModelBuilderMode.FLOAT,
                                                   append2output=self.all_interest_points)
@@ -135,12 +118,10 @@ class DistanceMetricCalculator(MetricCalculator):
         # Initiating baseline_tensors_list since it is not initiated in SensitivityEvaluationManager init.
         self.baseline_tensors_list = self._init_baseline_tensors_list()
 
-        # Computing Hessian-based scores for weighted average distance metric computation (only if requested),
-        # and assigning distance_weighting method accordingly.
+        # Hessian-based scores for weighted average distance metric computation
         self.interest_points_hessians = None
-        if self.mp_config.use_hessian_based_scores is True:
-            self.interest_points_hessians = self._compute_hessian_based_scores()
-            self.mp_config.distance_weighting_method = lambda d: self.interest_points_hessians
+        if self.mp_config.distance_weighting_method == MpDistanceWeighting.HESSIAN:
+            self.interest_points_hessians = self._compute_hessian_based_scores(hessian_info_service)
 
     def compute(self, mp_model) -> float:
         """
@@ -153,13 +134,11 @@ class DistanceMetricCalculator(MetricCalculator):
             Computed metric.
         """
         ipts_distances, out_pts_distances = self._compute_distance(mp_model)
-        sensitivity_metric = self._compute_mp_distance_measure(ipts_distances, out_pts_distances,
-                                                               self.mp_config.distance_weighting_method)
+        sensitivity_metric = self._compute_mp_distance_measure(ipts_distances, out_pts_distances)
         return sensitivity_metric
 
     def _init_metric_points_lists(self,
-                                  points: List[BaseNode],
-                                  norm_mse: bool = False) -> Tuple[List[Callable], List[int]]:
+                                  points: List[BaseNode]) -> Tuple[List[Callable], List[int]]:
         """
         Initiates required lists for future use when computing the sensitivity metric.
         Each point on which the metric is computed uses a dedicated distance function based on its type.
@@ -167,7 +146,6 @@ class DistanceMetricCalculator(MetricCalculator):
 
         Args:
             points: The set of nodes in the graph for which we need to initiate the lists.
-            norm_mse: whether to normalize mse distance function.
 
         Returns: A lists with distance functions and an axis list for each node.
 
@@ -177,7 +155,7 @@ class DistanceMetricCalculator(MetricCalculator):
         for n in points:
             distance_fn, axis = self.fw_impl.get_mp_node_distance_fn(n,
                                                                      compute_distance_fn=self.mp_config.compute_distance_fn,
-                                                                     norm_mse=norm_mse)
+                                                                     norm_mse=self.use_normalized_mse)
             distance_fns_list.append(distance_fn)
             # Axis is needed only for KL Divergence calculation, otherwise we use per-tensor computation
             axis_list.append(axis if distance_fn == compute_kl_divergence else None)
@@ -190,12 +168,14 @@ class DistanceMetricCalculator(MetricCalculator):
         return [self.fw_impl.to_numpy(self.fw_impl.sensitivity_eval_inference(self.ref_model, images))
                 for images in self.images_batches]
 
-    def _compute_hessian_based_scores(self) -> np.ndarray:
+    def _compute_hessian_based_scores(self, hessian_info_service: HessianInfoService) -> np.ndarray:
         """
         Compute Hessian-based scores for each interest point.
+        Args:
+            hessian_info_service: Hessian service.
 
-        Returns: A vector of scores, one for each interest point,
-         to be used for the distance metric weighted average computation.
+        Returns:
+            A vector of scores, one for each interest point, to be used for the distance metric weighted average computation.
 
         """
         # Create a request for Hessian approximation scores with specific configurations
@@ -209,12 +189,16 @@ class DistanceMetricCalculator(MetricCalculator):
                                                     n_samples=self.mp_config.num_of_images)
 
         # Fetch the Hessian approximation scores for the current interest point
-        nodes_approximations = self.hessian_info_service.fetch_hessian(request=hessian_info_request)
+        nodes_approximations = hessian_info_service.fetch_hessian(request=hessian_info_request)
         approx_by_image = np.stack([nodes_approximations[n.name] for n in self.interest_points],
                                    axis=1)  # samples X nodes
 
         # Return the mean approximation value across all images for each interest point
-        return np.mean(approx_by_image, axis=0)
+        scores = np.mean(approx_by_image, axis=0)
+        if scores.ndim == 2 and scores.shape[1] == 1:
+            scores = np.squeeze(scores, 1)
+        assert scores.ndim == 1, f'Expected a vector of hessians, got tensor of shape {scores.shape}'
+        return scores
 
     def _compute_points_distance(self,
                                  baseline_tensors: List[Any],
@@ -282,39 +266,58 @@ class DistanceMetricCalculator(MetricCalculator):
 
         return ipts_distances, out_pts_distances
 
-    @staticmethod
-    def _compute_mp_distance_measure(ipts_distances: np.ndarray,
-                                     out_pts_distances: np.ndarray,
-                                     metrics_weights_fn: Callable) -> float:
+    def _compute_mp_distance_measure(self, ipts_distances: np.ndarray, out_pts_distances: np.ndarray) -> float:
         """
         Computes the final distance value out of a distance matrix.
 
         Args:
-            ipts_distances: A matrix that contains the distances between the baseline and MP models
-                for each interest point.
-            out_pts_distances: A matrix that contains the distances between the baseline and MP models
-                for each output point.
-            metrics_weights_fn: A callable that produces the scores to compute weighted distance for interest points.
+            ipts_distances: A matrix that contains the distances between the reference and MP models
+                for each interest point, of shape (num interest points, num samples,).
+            out_pts_distances: A matrix that contains the distances between the reference and MP models
+                for each output point, of shape (num output points, num samples,).
 
-        Returns: Distance value.
+        Returns:
+            Distance value.
         """
-        mean_ipts_distance = 0
-        if len(ipts_distances) > 0:
-            mean_distance_per_layer = ipts_distances.mean(axis=1)
+        assert ipts_distances.size + out_pts_distances.size, 'Both interest and output points distances are empty.'
 
-            # Use weights such that every layer's distance is weighted differently (possibly).
-            weight_scores = metrics_weights_fn(ipts_distances)
-            weight_scores = np.asarray(weight_scores) if isinstance(weight_scores, List) else weight_scores
-            weight_scores = weight_scores.flatten()
+        ipts_metric = self._compute_ipts_distance_measure(ipts_distances) if ipts_distances.size else 0
 
-            mean_ipts_distance = np.average(mean_distance_per_layer, weights=weight_scores)
+        out_pts_metric = out_pts_distances.mean() if out_pts_distances.size else 0
 
-        mean_output_distance = 0
-        if len(out_pts_distances) > 0:
-            mean_distance_per_output = out_pts_distances.mean(axis=1)
-            mean_output_distance = np.average(mean_distance_per_output)
+        return ipts_metric + out_pts_metric
 
-        return mean_output_distance + mean_ipts_distance
+    def _compute_ipts_distance_measure(self, ipts_distances: np.ndarray) -> float:
+        """
+        Compute distance measure for interest points.
+
+        Args:
+            ipts_distances: a matrix of shape (num interest points, num samples,).
+
+        Returns:
+            Distance measure.
+        """
+        assert ipts_distances.ndim == 2, (f'Expected ipts_distances of shape shape (num interest points, num samples), '
+                                          f'got {ipts_distances.shape}')
+        method = self.mp_config.distance_weighting_method
+        if method == MpDistanceWeighting.AVG:
+            return ipts_distances.mean()
+        if method == MpDistanceWeighting.LAST_LAYER:
+            return ipts_distances[-1, :].mean()
+        if method == MpDistanceWeighting.HESSIAN:
+            return np.average(ipts_distances.mean(axis=1), weights=self.interest_points_hessians)
+        if method == MpDistanceWeighting.EXP:
+            assert self.use_normalized_mse
+            ipts_mean_distances = ipts_distances.mean(axis=1)
+            weights = 1 - np.exp(-ipts_mean_distances / self.mp_config.exp_distance_weighting_sigma)
+            if np.any(weights):
+                return np.average(ipts_mean_distances, weights=weights)
+            else:
+                Logger.warning('All weights for interest points are 0. If distances are very small, you might need to '
+                               'pass a smaller exp_distance_weighting_sigma.')
+                return 0
+
+        raise ValueError(f'Unexpected MpDistanceWeighting {method}')  # pragma: no cover
 
     def _get_images_batches(self, num_of_images: int) -> List[Any]:
         """

model_compression_toolkit/core/common/statistics_correction/compute_bias_correction_of_graph.py

@@ -115,7 +115,8 @@ def _compute_bias_correction(kernel: np.ndarray,
                              quantized_kernel: np.ndarray,
                              in_statistics_container: BaseStatsCollector,
                              output_channels_axis: int,
-                             input_channels_axis: int) -> Any:
+                             input_channels_axis: int,
+                             node_name: str) -> Any:
     """
     Compute the bias correction term for the bias in the error on the layer’s output,
     that is introduced by the weights quantization.
@@ -154,14 +155,14 @@ def _compute_bias_correction(kernel: np.ndarray,
 
     # Sanity validation
     if is_non_positive_integer(num_groups) or is_non_positive_integer(num_out_channels / num_groups):
-        Logger.warning("Skipping bias correction due to valiation problem.")
+        Logger.warning(f"Skipping bias correction due to validation problem in node {node_name}.")
         return correction_term
 
     num_out_channels_per_group = int(num_out_channels / num_groups)
 
     # In Pytorch the output of group conv is separated into respective groups is
     # viewed as follows: (batch, channel, ngroups, h, w),
-    # i.e each group is consistently viewed one after the other
+    # i.e. each group is consistently viewed one after the other
     # For an example, check out: https://discuss.pytorch.org/t/group-convolution-output-order/88258
     mu_split = np.split(mu, num_groups)
     eps_split = np.split(eps, num_groups, 0)
@@ -205,5 +206,6 @@ def _get_bias_correction_term_of_node(input_channels_axis: int,
                                           quantized_kernel,
                                           node_in_stats_collector,
                                           output_channels_axis,
-                                          input_channels_axis)
+                                          input_channels_axis,
+                                          n.name)
     return correction

model_compression_toolkit/core/common/mixed_precision/sensitivity_eval/distance_weighting.py

@@ -1,76 +0,0 @@
-# Copyright 2021 Sony Semiconductor Israel, Inc. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-from enum import Enum
-from functools import partial
-
-import numpy as np
-
-
-def get_average_weights(distance_matrix: np.ndarray) -> np.ndarray:
-    """
-    Get weights for weighting the sensitivity among different layers when evaluating MP configurations on
-    model's sensitivity. This function returns equal weights for each layer, such that the sensitivity
-    is averaged over all layers.
-
-    Args:
-        distance_matrix: Numpy array at shape (L,M): L -number of interest points, M number of samples.
-        The matrix contain the distance for each interest point at each sample.
-
-    Returns:
-        Numpy array containing equal weights for sensitivity weighting.
-    """
-
-    num_nodes = len(distance_matrix)
-    return np.asarray([1 / num_nodes for _ in range(num_nodes)])
-
-
-def get_last_layer_weights(distance_matrix: np.ndarray) -> np.ndarray:
-    """
-    Get weights for weighting the sensitivity among different layers when evaluating MP configurations on
-    model's sensitivity. This function returns weights for each layer, such that the sensitivity
-    is computed using only the last layer of the model (all other weights are zero).
-
-    Args:
-        distance_matrix: Numpy array at shape (L,M): L -number of interest points, M number of samples.
-        The matrix contain the distance for each interest point at each sample.
-
-    Returns:
-        Numpy array containing weights for sensitivity weighting (all zero but the last one).
-    """
-    num_nodes = len(distance_matrix)
-    w = np.asarray([0 for _ in range(num_nodes)])
-    w[-1] = 1
-    return w
-
-
-class MpDistanceWeighting(Enum):
-    """
-    Defines mixed precision distance metric weighting methods.
-    The enum values can be used to call a function on a set of arguments and key-arguments.
-
-     AVG - take the average distance on all computed layers.
-
-     LAST_LAYER - take only the distance of the last layer output.
-
-    """
-
-    AVG = partial(get_average_weights)
-    LAST_LAYER = partial(get_last_layer_weights)
-
-    def __call__(self, distance_matrix: np.ndarray) -> np.ndarray:
-        return self.value(distance_matrix)
-
-    def __deepcopy__(self, memo):
-        return self