mct-nightly 2.0.0.20240417.406__tar.gz → 2.0.0.20240418.439__tar.gz

{mct-nightly-2.0.0.20240417.406 → mct-nightly-2.0.0.20240418.439}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mct-nightly
-Version: 2.0.0.20240417.406
+Version: 2.0.0.20240418.439
 Summary: A Model Compression Toolkit for neural networks
 Home-page: UNKNOWN
 License: UNKNOWN

{mct-nightly-2.0.0.20240417.406 → mct-nightly-2.0.0.20240418.439}/mct_nightly.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mct-nightly
-Version: 2.0.0.20240417.406
+Version: 2.0.0.20240418.439
 Summary: A Model Compression Toolkit for neural networks
 Home-page: UNKNOWN
 License: UNKNOWN

{mct-nightly-2.0.0.20240417.406 → mct-nightly-2.0.0.20240418.439}/mct_nightly.egg-info/SOURCES.txt

@@ -144,6 +144,7 @@ model_compression_toolkit/core/common/substitutions/batchnorm_reconstruction.py
 model_compression_toolkit/core/common/substitutions/batchnorm_refusing.py
 model_compression_toolkit/core/common/substitutions/linear_collapsing.py
 model_compression_toolkit/core/common/substitutions/linear_collapsing_substitution.py
+model_compression_toolkit/core/common/substitutions/remove_identity.py
 model_compression_toolkit/core/common/substitutions/residual_collapsing.py
 model_compression_toolkit/core/common/substitutions/scale_equalization.py
 model_compression_toolkit/core/common/substitutions/shift_negative_activation.py
@@ -183,6 +184,7 @@ model_compression_toolkit/core/keras/graph_substitutions/substitutions/linear_co
 model_compression_toolkit/core/keras/graph_substitutions/substitutions/matmul_substitution.py
 model_compression_toolkit/core/keras/graph_substitutions/substitutions/multi_head_attention_decomposition.py
 model_compression_toolkit/core/keras/graph_substitutions/substitutions/relu_bound_to_power_of_2.py
+model_compression_toolkit/core/keras/graph_substitutions/substitutions/remove_identity.py
 model_compression_toolkit/core/keras/graph_substitutions/substitutions/residual_collapsing.py
 model_compression_toolkit/core/keras/graph_substitutions/substitutions/scale_equalization.py
 model_compression_toolkit/core/keras/graph_substitutions/substitutions/separableconv_decomposition.py
@@ -247,6 +249,7 @@ model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/linear_
 model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/multi_head_attention_decomposition.py
 model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/permute_call_method.py
 model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/relu_bound_to_power_of_2.py
+model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/remove_identity.py
 model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/reshape_with_static_shapes.py
 model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/residual_collapsing.py
 model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/scale_equalization.py

{mct-nightly-2.0.0.20240417.406 → mct-nightly-2.0.0.20240418.439}/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.0.0.20240417.000406"
+__version__ = "2.0.0.20240418.000439"

{mct-nightly-2.0.0.20240417.406 → mct-nightly-2.0.0.20240418.439}/model_compression_toolkit/constants.py

@@ -97,6 +97,8 @@ UPPER_FACTOR = 1.2
 DEC_RANGE_BOTTOM = 0.97
 DEC_RANGE_UPPER = 1.03
 
+NUM_QPARAM_HESSIAN_SAMPLES = 16
+
 # Resource utilization computation parameters
 BITS_TO_BYTES = 8.0
 

{mct-nightly-2.0.0.20240417.406 → mct-nightly-2.0.0.20240418.439}/model_compression_toolkit/core/common/graph/base_node.py

@@ -14,7 +14,7 @@
 # ==============================================================================
 
 import copy
-from typing import Dict, Any, Tuple, List, Type
+from typing import Dict, Any, Tuple, List, Type, Union
 
 import numpy as np
 

{mct-nightly-2.0.0.20240417.406 → mct-nightly-2.0.0.20240418.439}/model_compression_toolkit/core/common/hessian/hessian_info_service.py

@@ -17,7 +17,6 @@ from functools import partial
 from typing import Callable, List
 
 from model_compression_toolkit.constants import HESSIAN_NUM_ITERATIONS
-from model_compression_toolkit.core.common import Graph
 from model_compression_toolkit.core.common.hessian.trace_hessian_request import TraceHessianRequest
 from model_compression_toolkit.logger import Logger
 
@@ -38,7 +37,7 @@ class HessianInfoService:
     """
 
     def __init__(self,
-                 graph: Graph,
+                 graph,
                  representative_dataset: Callable,
                  fw_impl,
                  num_iterations_for_approximation: int = HESSIAN_NUM_ITERATIONS
@@ -151,7 +150,7 @@ class HessianInfoService:
         if required_size==0:
             return []
 
-        Logger.info(f"Ensuring {required_size} Hessian-trace approximation for node {trace_hessian_request.target_node}.")
+        Logger.info(f"\nEnsuring {required_size} Hessian-trace approximation for node {trace_hessian_request.target_node}.")
 
         # Replace request of a reused target node with a request of the 'reuse group'.
         if trace_hessian_request.target_node.reuse_group:

{mct-nightly-2.0.0.20240417.406 → mct-nightly-2.0.0.20240418.439}/model_compression_toolkit/core/common/hessian/trace_hessian_request.py

@@ -16,8 +16,6 @@ from typing import List
 
 from enum import Enum
 
-from model_compression_toolkit.core.common import BaseNode
-
 
 class HessianMode(Enum):
     """
@@ -54,7 +52,7 @@ class TraceHessianRequest:
     def __init__(self,
                  mode: HessianMode,
                  granularity: HessianInfoGranularity,
-                 target_node: BaseNode,
+                 target_node,
                  ):
         """
         Attributes:

{mct-nightly-2.0.0.20240417.406 → mct-nightly-2.0.0.20240418.439}/model_compression_toolkit/core/common/quantization/quantization_config.py

@@ -26,14 +26,16 @@ class QuantizationErrorMethod(Enum):
 
     NOCLIPPING - Use min/max values as thresholds.
 
-    MSE - Use min square error for minimizing quantization noise.
+    MSE - Use mean square error for minimizing quantization noise.
 
-    MAE - Use min absolute error for minimizing quantization noise.
+    MAE - Use mean absolute error for minimizing quantization noise.
 
     KL - Use KL-divergence to make signals distributions to be similar as possible.
 
     Lp - Use Lp-norm to minimizing quantization noise.
 
+    HMSE - Use Hessian-based mean squared error for minimizing quantization noise. This method is using Hessian scores to factorize more valuable parameters when computing the error induced by quantization.
+
     """
 
     NOCLIPPING = 0
@@ -41,6 +43,7 @@ class QuantizationErrorMethod(Enum):
     MAE = 2
     KL = 4
     LP = 5
+    HMSE = 6
 
 
 class QuantizationConfig:

{mct-nightly-2.0.0.20240417.406 → mct-nightly-2.0.0.20240418.439}/model_compression_toolkit/core/common/quantization/quantization_params_generation/error_functions.py

@@ -13,13 +13,16 @@
 # limitations under the License.
 # ==============================================================================
 from copy import deepcopy
-from typing import Tuple, Callable
+from typing import Tuple, Callable, List
 import numpy as np
 import model_compression_toolkit.core.common.quantization.quantization_config as qc
+from model_compression_toolkit.core.common.hessian import TraceHessianRequest, HessianMode, HessianInfoGranularity, \
+    HessianInfoService
 from model_compression_toolkit.core.common.similarity_analyzer import compute_mse, compute_mae, compute_lp_norm
 from model_compression_toolkit.target_platform_capabilities.target_platform import QuantizationMethod
-from model_compression_toolkit.constants import FLOAT_32
-from model_compression_toolkit.core.common.quantization.quantizers.quantizers_helpers import uniform_quantize_tensor
+from model_compression_toolkit.constants import FLOAT_32, NUM_QPARAM_HESSIAN_SAMPLES
+from model_compression_toolkit.core.common.quantization.quantizers.quantizers_helpers import uniform_quantize_tensor, \
+    reshape_tensor_for_per_channel_search
 
 
 def _mse_error_histogram(q_bins: np.ndarray,
@@ -371,13 +374,63 @@ def _get_sliced_histogram(bins: np.ndarray,
     return bins_subset, counts_subset
 
 
+def _compute_hessian_for_hmse(node,
+                              hessian_info_service: HessianInfoService,
+                              num_hessian_samples: int = NUM_QPARAM_HESSIAN_SAMPLES) -> List[np.ndarray]:
+    """
+    Compute and retrieve Hessian-based scores for using during HMSE error computation.
+
+    Args:
+        node: The node to compute Hessian-based scores for.
+        hessian_info_service: HessianInfoService object for retrieving Hessian-based scores.
+        num_hessian_samples: Number of samples to approximate Hessian-based scores on.
+
+    Returns: A list with computed Hessian-based scores tensors for the given node.
+
+    """
+    _request = TraceHessianRequest(mode=HessianMode.WEIGHTS,
+                                   granularity=HessianInfoGranularity.PER_ELEMENT,
+                                   target_node=node)
+    _scores_for_node = hessian_info_service.fetch_hessian(_request,
+                                                          required_size=num_hessian_samples)
+
+    return _scores_for_node
+
+
+def _hmse_error_function_wrapper(float_tensor: np.ndarray,
+                                 fxp_tensor: np.ndarray,
+                                 axis: int,
+                                 norm: bool,
+                                 hessian_scores: np.ndarray):
+    """
+    This function wraps the HMSE error method to enable using it during parameters selection.
+
+    Args:
+        float_tensor: Float tensor.
+        fxp_tensor: Quantized tensor.
+        axis: Axis along which the operation has been performed. If not None, then per-channel computation is expected.
+        norm: Indicates whether to normalize the result of the error function.
+        hessian_scores: A tensor with Hessian-based scores to use for Hessian-based MSE (HMSE) error computation.
+
+    Returns: The HMSE error between the float and fixed-point tensors.
+
+    """
+    if axis is not None:
+        hessian_scores = reshape_tensor_for_per_channel_search(hessian_scores, 0)
+
+    return compute_mse(float_tensor, fxp_tensor, axis, norm, weights=hessian_scores)
+
+
 def get_threshold_selection_tensor_error_function(quantization_method: QuantizationMethod,
                                                   quant_error_method: qc.QuantizationErrorMethod,
                                                   p: int,
                                                   axis: int = None,
                                                   norm: bool = False,
                                                   n_bits: int = 8,
-                                                  signed: bool = True) -> Callable:
+                                                  signed: bool = True,
+                                                  node=None,
+                                                  hessian_info_service: HessianInfoService = None,
+                                                  num_hessian_samples: int = NUM_QPARAM_HESSIAN_SAMPLES) -> Callable:
     """
     Returns the error function compatible to the provided threshold method,
     to be used in the threshold optimization search for tensor quantization.
@@ -389,6 +442,9 @@ def get_threshold_selection_tensor_error_function(quantization_method: Quantizat
         norm: Indicates whether to normalize the result of the error function.
         n_bits: Number of bits used to quantize the tensor.
         signed: Indicates whether the input is signed.
+        node: The node for which the quantization error is computed (used only with HMSE error method).
+        hessian_info_service: HessianInfoService object for retrieving Hessian-based scores (used only with HMSE error method).
+        num_hessian_samples: Number of samples to approximate Hessian-based scores on (used only with HMSE error method).
 
     Returns: a Callable method that calculates the error between a tensor and a quantized tensor.
     """
@@ -418,6 +474,13 @@ def get_threshold_selection_tensor_error_function(quantization_method: Quantizat
                                                                           n_bits=n_bits,
                                                                           per_channel=True)
 
+    if quant_error_method == qc.QuantizationErrorMethod.HMSE:
+        node_hessian_scores = _compute_hessian_for_hmse(node, hessian_info_service, num_hessian_samples)
+        node_hessian_scores = np.sqrt(np.mean(node_hessian_scores, axis=0))
+
+        return lambda x, y, threshold: _hmse_error_function_wrapper(x, y, norm=norm, axis=axis,
+                                                                    hessian_scores=node_hessian_scores)
+
     quant_method_error_function_mapping = {
         qc.QuantizationErrorMethod.MSE: lambda x, y, threshold: compute_mse(x, y, norm=norm, axis=axis),
         qc.QuantizationErrorMethod.MAE: lambda x, y, threshold: compute_mae(x, y, norm=norm, axis=axis),

{mct-nightly-2.0.0.20240417.406 → mct-nightly-2.0.0.20240418.439}/model_compression_toolkit/core/common/quantization/quantization_params_generation/lut_kmeans_params.py

@@ -18,7 +18,8 @@ from sklearn.cluster import KMeans
 
 import model_compression_toolkit.core.common.quantization.quantization_config as qc
 from model_compression_toolkit.constants import LUT_VALUES, MIN_THRESHOLD, SCALE_PER_CHANNEL, \
-    LUT_VALUES_BITWIDTH, THRESHOLD
+    LUT_VALUES_BITWIDTH, THRESHOLD, NUM_QPARAM_HESSIAN_SAMPLES
+from model_compression_toolkit.core.common.hessian import HessianInfoService
 from model_compression_toolkit.core.common.quantization.quantizers.quantizers_helpers import \
     max_power_of_two, int_quantization_with_threshold
 from model_compression_toolkit.core.common.quantization.quantization_params_generation.symmetric_selection import \
@@ -37,7 +38,10 @@ def lut_kmeans_tensor(tensor_data: np.ndarray,
                       n_iter: int = 10,
                       min_threshold: float = MIN_THRESHOLD,
                       quant_error_method: qc.QuantizationErrorMethod = None,
-                      is_symmetric=False) -> dict:
+                      is_symmetric=False,
+                      node=None,
+                      hessian_info_service: HessianInfoService = None,
+                      num_hessian_samples: int = NUM_QPARAM_HESSIAN_SAMPLES) -> dict:
     """
     The quantizer first finds the closest max value per channel of tensor_data.
     Now, we divide tensor_data with the threshold vector per channel. In addition, we scale the result to the range
@@ -53,7 +57,10 @@ def lut_kmeans_tensor(tensor_data: np.ndarray,
         n_iter: Number of iterations to search_methods for the optimal threshold.
         min_threshold: Minimal threshold to chose when the computed one is smaller.
         quant_error_method: an error function to optimize the parameters' selection accordingly (not used for this method).
-        is_symmetric (bool): Whether to apply symmetric weight quantization (default is False, meaning power of 2 quantization)
+        is_symmetric (bool): Whether to apply symmetric weight quantization (default is False, meaning power of 2 quantization).
+        node: The node for which the quantization error is computed (not used for this method).
+        hessian_info_service: HessianInfoService object for retrieving Hessian-based scores (not used for this method).
+        num_hessian_samples: Number of samples to approximate Hessian-based scores on (not used for this method).
 
     Returns:
         A dictionary containing the cluster assignments according to the k-means algorithm,

{mct-nightly-2.0.0.20240417.406 → mct-nightly-2.0.0.20240418.439}/model_compression_toolkit/core/common/quantization/quantization_params_generation/power_of_two_selection.py

@@ -15,7 +15,8 @@
 import numpy as np
 
 import model_compression_toolkit.core.common.quantization.quantization_config as qc
-from model_compression_toolkit.constants import MIN_THRESHOLD, THRESHOLD
+from model_compression_toolkit.constants import MIN_THRESHOLD, THRESHOLD, NUM_QPARAM_HESSIAN_SAMPLES
+from model_compression_toolkit.core.common.hessian import HessianInfoService
 from model_compression_toolkit.core.common.quantization.quantization_params_generation.qparams_search import \
     qparams_selection_tensor_search, qparams_selection_histogram_search
 from model_compression_toolkit.core.common.quantization.quantizers.quantizers_helpers import max_power_of_two, get_tensor_max
@@ -31,7 +32,11 @@ def power_of_two_selection_tensor(tensor_data: np.ndarray,
                                   channel_axis: int = 1,
                                   n_iter: int = 10,
                                   min_threshold: float = MIN_THRESHOLD,
-                                  quant_error_method: qc.QuantizationErrorMethod = qc.QuantizationErrorMethod.MSE) -> dict:
+                                  quant_error_method: qc.QuantizationErrorMethod = qc.QuantizationErrorMethod.MSE,
+                                  node=None,
+                                  hessian_info_service: HessianInfoService = None,
+                                  num_hessian_samples: int = NUM_QPARAM_HESSIAN_SAMPLES,
+                                  ) -> dict:
     """
     Compute the power of two threshold based on the provided QuantizationErrorMethod to quantize the tensor.
     Different search is applied, depends on the value of the selected QuantizationErrorMethod.
@@ -45,6 +50,9 @@ def power_of_two_selection_tensor(tensor_data: np.ndarray,
         n_iter: Number of iterations to search for the optimal threshold (not used for this method).
         min_threshold: Minimal threshold to use if threshold is too small (not used for this method).
         quant_error_method: an error function to optimize the parameters' selection accordingly.
+        node: The node for which the quantization error is computed (used only with HMSE error method).
+        hessian_info_service: HessianInfoService object for retrieving Hessian-based scores (used only with HMSE error method).
+        num_hessian_samples: Number of samples to approximate Hessian-based scores on (used only with HMSE error method).
 
     Returns:
         Power of two threshold to quantize the tensor in a power of 2 manner.
@@ -57,8 +65,10 @@ def power_of_two_selection_tensor(tensor_data: np.ndarray,
         signed = True  # weights are always signed
         axis = -1 if per_channel else None
         error_function = get_threshold_selection_tensor_error_function(QuantizationMethod.POWER_OF_TWO,
-                                                                       quant_error_method, p, axis=axis, norm=False, n_bits=n_bits,
-                                                                       signed=signed)
+                                                                       quant_error_method, p, axis=axis, norm=False,
+                                                                       n_bits=n_bits, signed=signed, node=node,
+                                                                       hessian_info_service=hessian_info_service,
+                                                                       num_hessian_samples=num_hessian_samples)
         threshold = qparams_selection_tensor_search(error_function,
                                                     tensor_data,
                                                     n_bits,

{mct-nightly-2.0.0.20240417.406 → mct-nightly-2.0.0.20240418.439}/model_compression_toolkit/core/common/quantization/quantization_params_generation/qparams_computation.py

@@ -12,10 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+import copy
+
 from tqdm import tqdm
 from typing import List
 
+from model_compression_toolkit.constants import NUM_QPARAM_HESSIAN_SAMPLES
+from model_compression_toolkit.core import QuantizationErrorMethod
 from model_compression_toolkit.core.common import Graph, BaseNode
+from model_compression_toolkit.core.common.hessian import HessianInfoService
 from model_compression_toolkit.core.common.quantization.quantization_params_generation.qparams_activations_computation \
     import get_activations_qparams
 from model_compression_toolkit.core.common.quantization.quantization_params_generation.qparams_weights_computation import \
@@ -25,7 +30,9 @@ from model_compression_toolkit.logger import Logger
 
 def calculate_quantization_params(graph: Graph,
                                   nodes: List[BaseNode] = [],
-                                  specific_nodes: bool = False):
+                                  specific_nodes: bool = False,
+                                  hessian_info_service: HessianInfoService = None,
+                                  num_hessian_samples: int = NUM_QPARAM_HESSIAN_SAMPLES):
     """
     For a graph, go over its nodes, compute quantization params (for both weights and activations according
     to the given framework info), and create and attach a NodeQuantizationConfig to each node (containing the
@@ -39,6 +46,8 @@ def calculate_quantization_params(graph: Graph,
         graph: Graph to compute its nodes' thresholds.
         nodes: List of nodes to compute their thresholds instead of computing it for all nodes in the graph.
         specific_nodes: Flag to compute thresholds for only specific nodes.
+        hessian_info_service: HessianInfoService object for retrieving Hessian-based scores (used only with HMSE error method).
+        num_hessian_samples: Number of samples to approximate Hessian-based scores on (used only with HMSE error method).
 
     """
 
@@ -60,10 +69,28 @@ def calculate_quantization_params(graph: Graph,
                         output_channels_axis = channels_axis[0]
                     else:
                         output_channels_axis = None
+
+                    mod_attr_cfg = attr_cfg
+
+                    if attr_cfg.weights_error_method == QuantizationErrorMethod.HMSE:
+                        kernel_attr_name = graph.fw_info.get_kernel_op_attributes(n.type)
+                        if len(kernel_attr_name) > 0:
+                            kernel_attr_name = kernel_attr_name[0]
+
+                        if kernel_attr_name is None or kernel_attr_name not in attr:
+                            Logger.warning(f"The HMSE error method for parameters selection is only supported for "
+                                           f"kernel weights attributes. Running parameters selection for attribute "
+                                           f"'{attr}' in node '{n.name}' with the default MSE error method instead.")
+                            mod_attr_cfg = copy.deepcopy(attr_cfg)
+                            mod_attr_cfg.weights_error_method = QuantizationErrorMethod.MSE
+
                     weights_params = get_weights_qparams(n.get_weights_by_keys(attr),
                                                          candidate_qc.weights_quantization_cfg,
-                                                         attr_cfg,
-                                                         output_channels_axis)
+                                                         mod_attr_cfg,
+                                                         output_channels_axis,
+                                                         node=n,
+                                                         hessian_info_service=hessian_info_service,
+                                                         num_hessian_samples=num_hessian_samples)
                     attr_cfg.set_weights_quantization_param(weights_params)
 
             if n.is_activation_quantization_enabled():

{mct-nightly-2.0.0.20240417.406 → mct-nightly-2.0.0.20240418.439}/model_compression_toolkit/core/common/quantization/quantization_params_generation/qparams_weights_computation.py

@@ -12,11 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from typing import Dict, Any, Tuple
+from typing import Dict, Any
 
 import numpy as np
 
-from model_compression_toolkit.logger import Logger
+from model_compression_toolkit.constants import NUM_QPARAM_HESSIAN_SAMPLES
+from model_compression_toolkit.core.common.hessian import HessianInfoService
 from model_compression_toolkit.defaultdict import DefaultDict
 from model_compression_toolkit.core.common.framework_info import FrameworkInfo
 from model_compression_toolkit.core.common.quantization.node_quantization_config import NodeWeightsQuantizationConfig, \
@@ -27,31 +28,40 @@ from model_compression_toolkit.core.common.quantization.node_quantization_config
 dummy_channel_mapping = DefaultDict(default_value=(None, None))
 
 
-def get_weights_qparams(kernel: np.ndarray,
+def get_weights_qparams(weights_attr_values: np.ndarray,
                         weights_quant_config: NodeWeightsQuantizationConfig,
                         attr_quant_config: WeightsAttrQuantizationConfig,
-                        output_channels_axis: int) -> Dict[Any, Any]:
+                        output_channels_axis: int,
+                        node=None,
+                        hessian_info_service: HessianInfoService = None,
+                        num_hessian_samples: int = NUM_QPARAM_HESSIAN_SAMPLES) -> Dict[Any, Any]:
     """
     Compute thresholds to quantize a kernel according to a NodeWeightsQuantizationConfig
     instance.
 
     Args:
-        kernel: Kernel to compute the quantization thresholds to.
+        weights_attr_values: Weights attribute parameter to compute the quantization thresholds for.
         weights_quant_config: Weights quantization configuration to define how the thresholds are computed.
         attr_quant_config: A specific weights attribute quantization configuration to get its params.
         output_channels_axis: Index of the kernel output channels dimension.
+        node: The node for which the quantization error is computed (used only with HMSE error method).
+        hessian_info_service: HessianInfoService object for retrieving Hessian-based scores (used only with HMSE error method).
+        num_hessian_samples: Number of samples to approximate Hessian-based scores on (used only with HMSE error method).
 
     Returns:
         A dictionary with the quantization threshold of the kernel.
     """
     if attr_quant_config.weights_quantization_params_fn is not None:
-        weights_params = attr_quant_config.weights_quantization_params_fn(kernel,
+        weights_params = attr_quant_config.weights_quantization_params_fn(weights_attr_values,
                                                                           p=attr_quant_config.l_p_value,
                                                                           n_bits=attr_quant_config.weights_n_bits,
                                                                           per_channel=attr_quant_config.weights_per_channel_threshold and output_channels_axis is not None,
                                                                           channel_axis=output_channels_axis,
                                                                           min_threshold=weights_quant_config.min_threshold,
-                                                                          quant_error_method=attr_quant_config.weights_error_method)
+                                                                          quant_error_method=attr_quant_config.weights_error_method,
+                                                                          node=node,
+                                                                          hessian_info_service=hessian_info_service,
+                                                                          num_hessian_samples=num_hessian_samples)
     else:
         weights_params = {}
 

{mct-nightly-2.0.0.20240417.406 → mct-nightly-2.0.0.20240418.439}/model_compression_toolkit/core/common/quantization/quantization_params_generation/symmetric_selection.py

@@ -15,7 +15,8 @@
 import numpy as np
 
 import model_compression_toolkit.core.common.quantization.quantization_config as qc
-from model_compression_toolkit.constants import MIN_THRESHOLD, THRESHOLD
+from model_compression_toolkit.constants import MIN_THRESHOLD, THRESHOLD, NUM_QPARAM_HESSIAN_SAMPLES
+from model_compression_toolkit.core.common.hessian import HessianInfoService
 from model_compression_toolkit.core.common.quantization.quantization_params_generation.error_functions import \
     get_threshold_selection_tensor_error_function, get_threshold_selection_histogram_error_function, _kl_error_histogram
 from model_compression_toolkit.core.common.quantization.quantization_params_generation.qparams_search import \
@@ -33,7 +34,10 @@ def symmetric_selection_tensor(tensor_data: np.ndarray,
                                channel_axis: int = 1,
                                n_iter: int = 10,
                                min_threshold: float = MIN_THRESHOLD,
-                               quant_error_method: qc.QuantizationErrorMethod = qc.QuantizationErrorMethod.MSE) -> dict:
+                               quant_error_method: qc.QuantizationErrorMethod = qc.QuantizationErrorMethod.MSE,
+                               node=None,
+                               hessian_info_service: HessianInfoService = None,
+                               num_hessian_samples: int = NUM_QPARAM_HESSIAN_SAMPLES) -> dict:
     """
     Compute the optimal threshold based on the provided QuantizationErrorMethod to quantize the tensor.
     Different search is applied, depends on the value of the selected QuantizationErrorMethod.
@@ -47,6 +51,9 @@ def symmetric_selection_tensor(tensor_data: np.ndarray,
         n_iter: Number of iterations to search for the optimal threshold (not used for this method).
         min_threshold: Minimal threshold to use if threshold is too small (not used for this method).
         quant_error_method: an error function to optimize the parameters' selection accordingly.
+        node: The node for which the quantization error is computed (used only with HMSE error method).
+        hessian_info_service: HessianInfoService object for retrieving Hessian-based scores (used only with HMSE error method).
+        num_hessian_samples: Number of samples to approximate Hessian-based scores on (used only with HMSE error method).
 
     Returns:
         Optimal threshold to quantize the tensor in a symmetric manner.
@@ -59,7 +66,11 @@ def symmetric_selection_tensor(tensor_data: np.ndarray,
     else:
         signed = True  # weights are always signed
         axis = -1 if per_channel else None
-        error_function = get_threshold_selection_tensor_error_function(QuantizationMethod.SYMMETRIC, quant_error_method, p, axis=axis, norm=False, n_bits=n_bits, signed=signed)
+        error_function = get_threshold_selection_tensor_error_function(QuantizationMethod.SYMMETRIC, quant_error_method,
+                                                                       p, axis=axis, norm=False, n_bits=n_bits,
+                                                                       signed=signed, node=node,
+                                                                       hessian_info_service=hessian_info_service,
+                                                                       num_hessian_samples=num_hessian_samples)
         threshold = qparams_symmetric_selection_tensor_search(error_function,
                                                               tensor_data,
                                                               tensor_max,

{mct-nightly-2.0.0.20240417.406 → mct-nightly-2.0.0.20240418.439}/model_compression_toolkit/core/common/quantization/quantization_params_generation/uniform_selection.py

@@ -15,7 +15,8 @@
 import numpy as np
 
 import model_compression_toolkit.core.common.quantization.quantization_config as qc
-from model_compression_toolkit.constants import MIN_THRESHOLD, RANGE_MIN, RANGE_MAX
+from model_compression_toolkit.constants import MIN_THRESHOLD, RANGE_MIN, RANGE_MAX, NUM_QPARAM_HESSIAN_SAMPLES
+from model_compression_toolkit.core.common.hessian import HessianInfoService
 from model_compression_toolkit.core.common.quantization.quantization_params_generation.qparams_search import \
     qparams_uniform_selection_tensor_search, qparams_uniform_selection_histogram_search
 from model_compression_toolkit.core.common.quantization.quantization_params_generation.error_functions import \
@@ -31,7 +32,10 @@ def uniform_selection_tensor(tensor_data: np.ndarray,
                              channel_axis: int = 1,
                              n_iter: int = 10,
                              min_threshold: float = MIN_THRESHOLD,
-                             quant_error_method: qc.QuantizationErrorMethod = qc.QuantizationErrorMethod.MSE) -> dict:
+                             quant_error_method: qc.QuantizationErrorMethod = qc.QuantizationErrorMethod.MSE,
+                             node=None,
+                             hessian_info_service: HessianInfoService = None,
+                             num_hessian_samples: int = NUM_QPARAM_HESSIAN_SAMPLES) -> dict:
     """
     Compute the optimal quantization range based on the provided QuantizationErrorMethod
     to uniformly quantize the tensor.
@@ -46,6 +50,9 @@ def uniform_selection_tensor(tensor_data: np.ndarray,
         n_iter: Number of iterations to search for the optimal threshold (not used for this method).
         min_threshold: Minimal threshold to use if threshold is too small (not used for this method).
         quant_error_method: an error function to optimize the range parameters' selection accordingly.
+        node: The node for which the quantization error is computed (used only with HMSE error method).
+        hessian_info_service: HessianInfoService object for retrieving Hessian-based scores (used only with HMSE error method).
+        num_hessian_samples: Number of samples to approximate Hessian-based scores on (used only with HMSE error method).
 
     Returns:
         Optimal quantization range to quantize the tensor uniformly.
@@ -57,7 +64,10 @@ def uniform_selection_tensor(tensor_data: np.ndarray,
         mm = tensor_min, tensor_max
     else:
         axis = -1 if per_channel else None
-        error_function = get_threshold_selection_tensor_error_function(QuantizationMethod.UNIFORM, quant_error_method, p, axis=axis, norm=False)
+        error_function = get_threshold_selection_tensor_error_function(QuantizationMethod.UNIFORM, quant_error_method,
+                                                                       p, axis=axis, norm=False, node=node,
+                                                                       hessian_info_service=hessian_info_service,
+                                                                       num_hessian_samples=num_hessian_samples)
         mm = qparams_uniform_selection_tensor_search(error_function,
                                                      tensor_data,
                                                      tensor_min,

{mct-nightly-2.0.0.20240417.406 → mct-nightly-2.0.0.20240418.439}/model_compression_toolkit/core/common/quantization/set_node_quantization_config.py

@@ -24,7 +24,8 @@ from model_compression_toolkit.core.common.graph.base_graph import Graph
 from model_compression_toolkit.core.common.quantization.candidate_node_quantization_config import \
     CandidateNodeQuantizationConfig
 from model_compression_toolkit.core.common.quantization.node_quantization_config import NodeActivationQuantizationConfig
-from model_compression_toolkit.core.common.quantization.quantization_config import QuantizationConfig
+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 \
@@ -36,19 +37,31 @@ from model_compression_toolkit.target_platform_capabilities.target_platform.op_q
 
 def set_quantization_configuration_to_graph(graph: Graph,
                                             quant_config: QuantizationConfig,
-                                            mixed_precision_enable: bool = False) -> Graph:
+                                            mixed_precision_enable: bool = False,
+                                            running_gptq: bool = False) -> Graph:
     """
     Add quantization configuration for each graph node.
 
     Args:
         graph: Graph for which to add quantization info to each node.
         quant_config: Quantization configuration containing parameters for how the graph should be quantized.
-        mixed_precision_enable: is mixed precision enabled
+        mixed_precision_enable: is mixed precision enabled.
+        running_gptq: Whether or not a GPTQ optimization is planned to run after the PTQ process.
 
     Returns:
         The graph with quantization configurations attached to each node in it.
     """
 
+    if quant_config.weights_error_method == QuantizationErrorMethod.HMSE:
+        if not running_gptq:
+            Logger.warning(f"The HMSE error method for parameters selection is only supported when running GPTQ "
+                           f"optimization due to long execution time that is not suitable for basic PTQ. "
+                           f"Using the default MSE error method instead.")
+            quant_config.weights_error_method = QuantizationErrorMethod.MSE
+        else:
+            Logger.warning("Using the HMSE error method for weights quantization parameters search. "
+                           "Note: This method may significantly increase runtime during the parameter search process.")
+
     for n in graph.nodes:
         set_quantization_configs_to_node(node=n,
                                          quant_config=quant_config,

{mct-nightly-2.0.0.20240417.406 → mct-nightly-2.0.0.20240418.439}/model_compression_toolkit/core/common/similarity_analyzer.py

@@ -18,6 +18,8 @@ from typing import Any
 import numpy as np
 
 from model_compression_toolkit.constants import EPS
+from model_compression_toolkit.logger import Logger
+
 
 #########################
 #  Helpful functions
@@ -87,7 +89,8 @@ def compute_mse(float_tensor: np.ndarray,
                 norm: bool = False,
                 norm_eps: float = 1e-8,
                 batch: bool = False,
-                axis: int = None) -> float:
+                axis: int = None,
+                weights: np.ndarray = None) -> float:
     """
     Compute the mean square error between two numpy arrays.
 
@@ -98,6 +101,7 @@ def compute_mse(float_tensor: np.ndarray,
         norm_eps: epsilon value for error normalization stability.
         batch: Whether to run batch similarity analysis or not.
         axis: Axis along which the operator has been computed.
+        weights: Weights tensor to use for computing Weighted-MSE error computation.
 
     Returns:
         The MSE distance between the two tensors.
@@ -107,7 +111,15 @@ def compute_mse(float_tensor: np.ndarray,
     float_flat = flatten_tensor(float_tensor, batch, axis)
     fxp_flat = flatten_tensor(fxp_tensor, batch, axis)
 
-    error = ((float_flat - fxp_flat) ** 2).mean(axis=-1)
+    if weights is not None:
+        w_flat = flatten_tensor(weights, batch, axis)
+        if w_flat.shape != float_flat.shape:
+            Logger.critical(f"Shape mismatch: The shape of the weights tensor {weights.shape} does not match the shape "
+                            f"of the input tensors {float_flat.shape} for Weighted-MSE computation.")  # pragma: no cover
+        error = ((w_flat * (float_flat - fxp_flat)) ** 2).mean(axis=-1)
+    else:
+        error = ((float_flat - fxp_flat) ** 2).mean(axis=-1)
+
     if norm:
         error /= ((float_flat ** 2).mean(axis=-1) + norm_eps)