mct-nightly 2.1.0.20240708.453__py3-none-any.whl → 2.1.0.20240710.440__py3-none-any.whl

model_compression_toolkit/core/keras/keras_implementation.py

@@ -21,11 +21,11 @@ from mct_quantizers import KerasQuantizationWrapper, KerasActivationQuantization
 from tensorflow.keras.models import Model
 
 from model_compression_toolkit.constants import HESSIAN_NUM_ITERATIONS
-from model_compression_toolkit.core.common.hessian import TraceHessianRequest, HessianMode, HessianInfoService
+from model_compression_toolkit.core.common.hessian import HessianScoresRequest, HessianMode, HessianInfoService
 from model_compression_toolkit.core.keras.graph_substitutions.substitutions.remove_identity import RemoveIdentity
-from model_compression_toolkit.core.keras.hessian.activation_trace_hessian_calculator_keras import \
-    ActivationTraceHessianCalculatorKeras
-from model_compression_toolkit.core.keras.hessian.weights_trace_hessian_calculator_keras import WeightsTraceHessianCalculatorKeras
+from model_compression_toolkit.core.keras.hessian.activation_hessian_scores_calculator_keras import \
+    ActivationHessianScoresCalculatorKeras
+from model_compression_toolkit.core.keras.hessian.weights_hessian_scores_calculator_keras import WeightsHessianScoresCalculatorKeras
 from model_compression_toolkit.exporter.model_wrapper.fw_agnostic.get_inferable_quantizers import \
     get_inferable_quantizers
 from model_compression_toolkit.exporter.model_wrapper.keras.builder.node_to_quantizer import \
@@ -364,7 +364,7 @@ class KerasImplementation(FrameworkImplementation):
             representative_data_gen: Dataset to use for retrieving images for the models inputs.
             fw_info: FrameworkInfo object with information about the specific framework's model.
             disable_activation_for_metric: Whether to disable activation quantization when computing the MP metric.
-            hessian_info_service: HessianInfoService to fetch approximations of the hessian traces for the float model.
+            hessian_info_service: HessianScoresService to fetch scores based on a Hessian-approximation for the float model.
 
         Returns:
             A SensitivityEvaluation object.
@@ -460,36 +460,36 @@ class KerasImplementation(FrameworkImplementation):
             return compute_cs
         return partial(compute_mse, norm=norm_mse)
 
-    def get_trace_hessian_calculator(self,
-                                     graph: Graph,
-                                     input_images: List[Any],
-                                     trace_hessian_request: TraceHessianRequest,
-                                     num_iterations_for_approximation: int = HESSIAN_NUM_ITERATIONS):
+    def get_hessian_scores_calculator(self,
+                                      graph: Graph,
+                                      input_images: List[Any],
+                                      hessian_scores_request: HessianScoresRequest,
+                                      num_iterations_for_approximation: int = HESSIAN_NUM_ITERATIONS):
         """
-        Get Keras trace hessian approximations calculator based on the trace hessian request.
+        Get Keras Hessian-approximation scores calculator based on the request.
         Args:
             input_images: Images to use for computation.
             graph: Float graph to compute the approximation of its different nodes.
-            trace_hessian_request: TraceHessianRequest to search for the desired calculator.
-            num_iterations_for_approximation: Number of iterations to use when approximating the Hessian trace.
-
-        Returns: TraceHessianCalculatorKeras to use for the trace hessian approximation computation for this request.
-
-        """
-        if trace_hessian_request.mode == HessianMode.ACTIVATION:
-            return ActivationTraceHessianCalculatorKeras(graph=graph,
-                                                         trace_hessian_request=trace_hessian_request,
-                                                         input_images=input_images,
-                                                         fw_impl=self,
-                                                         num_iterations_for_approximation=num_iterations_for_approximation)
-        elif trace_hessian_request.mode == HessianMode.WEIGHTS:
-            return WeightsTraceHessianCalculatorKeras(graph=graph,
-                                                      trace_hessian_request=trace_hessian_request,
-                                                      input_images=input_images,
-                                                      fw_impl=self,
-                                                      num_iterations_for_approximation=num_iterations_for_approximation)
+            hessian_scores_request: HessianScoresRequest to search for the desired calculator.
+            num_iterations_for_approximation: Number of iterations to use when approximating the Hessian scores.
+
+        Returns: HessianScoresCalculatorKeras to use for the Hessian-approximation scores computation for this request.
+
+        """
+        if hessian_scores_request.mode == HessianMode.ACTIVATION:
+            return ActivationHessianScoresCalculatorKeras(graph=graph,
+                                                          hessian_scores_request=hessian_scores_request,
+                                                          input_images=input_images,
+                                                          fw_impl=self,
+                                                          num_iterations_for_approximation=num_iterations_for_approximation)
+        elif hessian_scores_request.mode == HessianMode.WEIGHTS:
+            return WeightsHessianScoresCalculatorKeras(graph=graph,
+                                                       hessian_scores_request=hessian_scores_request,
+                                                       input_images=input_images,
+                                                       fw_impl=self,
+                                                       num_iterations_for_approximation=num_iterations_for_approximation)
         else:
-            Logger.critical(f"Unsupported Hessian mode for Keras: {trace_hessian_request.mode}.")   # pragma: no cover
+            Logger.critical(f"Unsupported Hessian mode for Keras: {hessian_scores_request.mode}.")   # pragma: no cover
 
     def is_output_node_compatible_for_hessian_score_computation(self,
                                                                 node: BaseNode) -> Any:

model_compression_toolkit/core/pytorch/hessian/{activation_trace_hessian_calculator_pytorch.py → activation_hessian_scores_calculator_pytorch.py}

@@ -21,48 +21,48 @@ import numpy as np
 
 from model_compression_toolkit.constants import MIN_HESSIAN_ITER, HESSIAN_COMP_TOLERANCE, HESSIAN_NUM_ITERATIONS
 from model_compression_toolkit.core.common import Graph
-from model_compression_toolkit.core.common.hessian import TraceHessianRequest, HessianInfoGranularity
+from model_compression_toolkit.core.common.hessian import HessianScoresRequest, HessianScoresGranularity
 from model_compression_toolkit.core.pytorch.back2framework.float_model_builder import FloatPyTorchModelBuilder
-from model_compression_toolkit.core.pytorch.hessian.trace_hessian_calculator_pytorch import \
-    TraceHessianCalculatorPytorch
+from model_compression_toolkit.core.pytorch.hessian.hessian_scores_calculator_pytorch import \
+    HessianScoresCalculatorPytorch
 from model_compression_toolkit.core.pytorch.utils import torch_tensor_to_numpy
 from model_compression_toolkit.logger import Logger
 import torch
 
 
-class ActivationTraceHessianCalculatorPytorch(TraceHessianCalculatorPytorch):
+class ActivationHessianScoresCalculatorPytorch(HessianScoresCalculatorPytorch):
     """
-    Pytorch implementation of the Trace Hessian approximation Calculator for activations.
+    Pytorch implementation of the Hessian approximation scores Calculator for activations.
     """
     def __init__(self,
                  graph: Graph,
                  input_images: List[torch.Tensor],
                  fw_impl,
-                 trace_hessian_request: TraceHessianRequest,
+                 hessian_scores_request: HessianScoresRequest,
                  num_iterations_for_approximation: int = HESSIAN_NUM_ITERATIONS):
         """
         Args:
             graph: Computational graph for the float model.
             input_images: List of input images for the computation.
-            fw_impl: Framework-specific implementation for trace Hessian approximation computation.
-            trace_hessian_request: Configuration request for which to compute the trace Hessian approximation.
-            num_iterations_for_approximation: Number of iterations to use when approximating the Hessian trace.
+            fw_impl: Framework-specific implementation for Hessian approximation scores computation.
+            hessian_scores_request: Configuration request for which to compute the Hessian approximation scores.
+            num_iterations_for_approximation: Number of iterations to use when approximating the Hessian scores.
 
         """
-        super(ActivationTraceHessianCalculatorPytorch, self).__init__(graph=graph,
-                                                                      input_images=input_images,
-                                                                      fw_impl=fw_impl,
-                                                                      trace_hessian_request=trace_hessian_request,
-                                                                      num_iterations_for_approximation=num_iterations_for_approximation)
+        super(ActivationHessianScoresCalculatorPytorch, self).__init__(graph=graph,
+                                                                       input_images=input_images,
+                                                                       fw_impl=fw_impl,
+                                                                       hessian_scores_request=hessian_scores_request,
+                                                                       num_iterations_for_approximation=num_iterations_for_approximation)
 
     def compute(self) -> List[np.ndarray]:
         """
-        Compute the approximation of the trace of the Hessian w.r.t the requested target nodes' activations.
+        Compute the scores that are based on the approximation of the Hessian w.r.t the requested target nodes' activations.
 
         Returns:
-            List[np.ndarray]: Approximated trace of the Hessian for the requested nodes.
+            List[np.ndarray]: Scores based on the approximated Hessian for the requested nodes.
         """
-        if self.hessian_request.granularity == HessianInfoGranularity.PER_TENSOR:
+        if self.hessian_request.granularity == HessianScoresGranularity.PER_TENSOR:
 
             model_output_nodes = [ot.node for ot in self.graph.get_outputs()]
 
@@ -99,7 +99,7 @@ class ActivationTraceHessianCalculatorPytorch(TraceHessianCalculatorPytorch):
             # First, we need to unfold all outputs that are given as list, to extract the actual output tensors
             output = self.concat_tensors(output_tensors)
 
-            ipts_hessian_trace_approx = [torch.tensor([0.0],
+            ipts_hessian_approx_scores = [torch.tensor([0.0],
                                                       requires_grad=True,
                                                       device=device)
                                          for _ in range(len(target_activation_tensors))]
@@ -109,7 +109,7 @@ class ActivationTraceHessianCalculatorPytorch(TraceHessianCalculatorPytorch):
                 v = torch.randn(output.shape, device=device)
                 f_v = torch.sum(v * output)
                 for i, ipt_tensor in enumerate(target_activation_tensors):  # Per Interest point activation tensor
-                    # Computing the hessian trace approximation by getting the gradient of (output * v)
+                    # Computing the hessian-approximation scores by getting the gradient of (output * v)
                     hess_v = autograd.grad(outputs=f_v,
                                            inputs=ipt_tensor,
                                            retain_graph=True,
@@ -121,24 +121,24 @@ class ActivationTraceHessianCalculatorPytorch(TraceHessianCalculatorPytorch):
                         continue  # pragma: no cover
 
                     # Mean over all dims but the batch (CXHXW for conv)
-                    hessian_trace_approx = torch.sum(hess_v ** 2.0, dim=tuple(d for d in range(1, len(hess_v.shape))))
+                    hessian_approx_scores = torch.sum(hess_v ** 2.0, dim=tuple(d for d in range(1, len(hess_v.shape))))
 
                     # Update node Hessian approximation mean over random iterations
-                    ipts_hessian_trace_approx[i] = (j * ipts_hessian_trace_approx[i] + hessian_trace_approx) / (j + 1)
+                    ipts_hessian_approx_scores[i] = (j * ipts_hessian_approx_scores[i] + hessian_approx_scores) / (j + 1)
 
                 # If the change to the maximal mean Hessian approximation is insignificant we stop the calculation
                 if j > MIN_HESSIAN_ITER:
                     if prev_mean_results is not None:
-                        new_mean_res = torch.mean(torch.stack(ipts_hessian_trace_approx), dim=1)
+                        new_mean_res = torch.mean(torch.stack(ipts_hessian_approx_scores), dim=1)
                         relative_delta_per_node = (torch.abs(new_mean_res - prev_mean_results) /
                                                    (torch.abs(new_mean_res) + 1e-6))
                         max_delta = torch.max(relative_delta_per_node)
                         if max_delta < HESSIAN_COMP_TOLERANCE:
                             break
-                prev_mean_results = torch.mean(torch.stack(ipts_hessian_trace_approx), dim=1)
+                prev_mean_results = torch.mean(torch.stack(ipts_hessian_approx_scores), dim=1)
 
             # Convert results to list of numpy arrays
-            hessian_results = [torch_tensor_to_numpy(h) for h in ipts_hessian_trace_approx]
+            hessian_results = [torch_tensor_to_numpy(h) for h in ipts_hessian_approx_scores]
             # Extend the Hessian tensors shape to align with expected return type
             # TODO: currently, only per-tensor Hessian is available for activation.
             #  Once implementing per-channel or per-element, this alignment needs to be verified and handled separately.
@@ -147,6 +147,6 @@ class ActivationTraceHessianCalculatorPytorch(TraceHessianCalculatorPytorch):
             return hessian_results
 
         else:  # pragma: no cover
-            Logger.critical(f"PyTorch activation Hessian's trace approximation does not support "
+            Logger.critical(f"PyTorch activation Hessian's approximation scores does not support "
                             f"{self.hessian_request.granularity} granularity.")
 

model_compression_toolkit/core/pytorch/hessian/{trace_hessian_calculator_pytorch.py → hessian_scores_calculator_pytorch.py}

@@ -17,38 +17,38 @@ from typing import Union, 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 import TraceHessianRequest
-from model_compression_toolkit.core.common.hessian.trace_hessian_calculator import TraceHessianCalculator
+from model_compression_toolkit.core.common.hessian import HessianScoresRequest
+from model_compression_toolkit.core.common.hessian.hessian_scores_calculator import HessianScoresCalculator
 from model_compression_toolkit.logger import Logger
 import torch
 
 
-class TraceHessianCalculatorPytorch(TraceHessianCalculator):
+class HessianScoresCalculatorPytorch(HessianScoresCalculator):
     """
-    Pytorch-specific implementation of the Trace Hessian approximation Calculator.
-    This class serves as a base for other Pytorch-specific trace Hessian approximation calculators.
+    Pytorch-specific implementation of the Hessian approximation scores Calculator.
+    This class serves as a base for other Pytorch-specific Hessian approximation scores calculators.
     """
     def __init__(self,
                  graph: Graph,
                  input_images: List[torch.Tensor],
                  fw_impl,
-                 trace_hessian_request: TraceHessianRequest,
+                 hessian_scores_request: HessianScoresRequest,
                  num_iterations_for_approximation: int = HESSIAN_NUM_ITERATIONS):
         """
 
         Args:
             graph: Computational graph for the float model.
             input_images: List of input images for the computation.
-            fw_impl: Framework-specific implementation for trace Hessian computation.
-            trace_hessian_request: Configuration request for which to compute the trace Hessian approximation.
-            num_iterations_for_approximation: Number of iterations to use when approximating the Hessian trace.
+            fw_impl: Framework-specific implementation for Hessian scores computation.
+            hessian_scores_request: Configuration request for which to compute the Hessian approximation scores.
+            num_iterations_for_approximation: Number of iterations to use when approximating the Hessian based scores.
 
         """
-        super(TraceHessianCalculatorPytorch, self).__init__(graph=graph,
-                                                            input_images=input_images,
-                                                            fw_impl=fw_impl,
-                                                            trace_hessian_request=trace_hessian_request,
-                                                            num_iterations_for_approximation=num_iterations_for_approximation)
+        super(HessianScoresCalculatorPytorch, self).__init__(graph=graph,
+                                                             input_images=input_images,
+                                                             fw_impl=fw_impl,
+                                                             hessian_scores_request=hessian_scores_request,
+                                                             num_iterations_for_approximation=num_iterations_for_approximation)
 
 
     def concat_tensors(self, tensors_to_concate: Union[torch.Tensor, List[torch.Tensor]]) -> torch.Tensor:

model_compression_toolkit/core/pytorch/hessian/{weights_trace_hessian_calculator_pytorch.py → weights_hessian_scores_calculator_pytorch.py}

@@ -18,41 +18,41 @@ import torch
 from torch import autograd
 import numpy as np
 from model_compression_toolkit.core.common import Graph
-from model_compression_toolkit.core.common.hessian import TraceHessianRequest, HessianInfoGranularity
-from model_compression_toolkit.core.pytorch.hessian.trace_hessian_calculator_pytorch import \
-    TraceHessianCalculatorPytorch
+from model_compression_toolkit.core.common.hessian import HessianScoresRequest, HessianScoresGranularity
+from model_compression_toolkit.core.pytorch.hessian.hessian_scores_calculator_pytorch import \
+    HessianScoresCalculatorPytorch
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.core.pytorch.back2framework.float_model_builder import FloatPyTorchModelBuilder
 from model_compression_toolkit.core.pytorch.default_framework_info import DEFAULT_PYTORCH_INFO
 from model_compression_toolkit.constants import HESSIAN_NUM_ITERATIONS, MIN_HESSIAN_ITER, HESSIAN_COMP_TOLERANCE, HESSIAN_EPS
 
 
-class WeightsTraceHessianCalculatorPytorch(TraceHessianCalculatorPytorch):
+class WeightsHessianScoresCalculatorPytorch(HessianScoresCalculatorPytorch):
     """
-    Pytorch-specific implementation of the Trace Hessian approximation computation w.r.t node's weights.
+    Pytorch-specific implementation of the Hessian approximation scores computation w.r.t node's weights.
     """
 
     def __init__(self,
                  graph: Graph,
                  input_images: List[torch.Tensor],
                  fw_impl,
-                 trace_hessian_request: TraceHessianRequest,
+                 hessian_scores_request: HessianScoresRequest,
                  num_iterations_for_approximation: int = HESSIAN_NUM_ITERATIONS):
         """
 
         Args:
             graph: Computational graph for the float model.
             input_images: List of input images for the computation.
-            fw_impl: Framework-specific implementation for trace Hessian computation.
-            trace_hessian_request: Configuration request for which to compute the trace Hessian approximation.
-            num_iterations_for_approximation: Number of iterations to use when approximating the Hessian trace.
+            fw_impl: Framework-specific implementation for Hessian scores computation.
+            hessian_scores_request: Configuration request for which to compute the Hessian approximation scores.
+            num_iterations_for_approximation: Number of iterations to use when approximating the Hessian scores.
         """
 
-        super(WeightsTraceHessianCalculatorPytorch, self).__init__(graph=graph,
-                                                                   input_images=input_images,
-                                                                   fw_impl=fw_impl,
-                                                                   trace_hessian_request=trace_hessian_request,
-                                                                   num_iterations_for_approximation=num_iterations_for_approximation)
+        super(WeightsHessianScoresCalculatorPytorch, self).__init__(graph=graph,
+                                                                    input_images=input_images,
+                                                                    fw_impl=fw_impl,
+                                                                    hessian_scores_request=hessian_scores_request,
+                                                                    num_iterations_for_approximation=num_iterations_for_approximation)
 
     def compute(self) -> List[np.ndarray]:
         """
@@ -77,9 +77,9 @@ class WeightsTraceHessianCalculatorPytorch(TraceHessianCalculatorPytorch):
         output_tensor = self.concat_tensors(outputs)
         device = output_tensor.device
 
-        ipts_hessian_trace_approx = [torch.tensor([0.0],
-                                                  requires_grad=True,
-                                                  device=device)
+        ipts_hessian_approx_scores = [torch.tensor([0.0],
+                                                   requires_grad=True,
+                                                   device=device)
                                      for _ in range(len(self.hessian_request.target_nodes))]
 
         prev_mean_results = None
@@ -107,9 +107,9 @@ class WeightsTraceHessianCalculatorPytorch(TraceHessianCalculatorPytorch):
                 # Get the output channel index
                 output_channel_axis, _ = DEFAULT_PYTORCH_INFO.kernel_channels_mapping.get(ipt_node.type)
                 shape_channel_axis = [i for i in range(len(weights_tensor.shape))]
-                if self.hessian_request.granularity == HessianInfoGranularity.PER_OUTPUT_CHANNEL:
+                if self.hessian_request.granularity == HessianScoresGranularity.PER_OUTPUT_CHANNEL:
                     shape_channel_axis.remove(output_channel_axis)
-                elif self.hessian_request.granularity == HessianInfoGranularity.PER_ELEMENT:
+                elif self.hessian_request.granularity == HessianScoresGranularity.PER_ELEMENT:
                     shape_channel_axis = ()
 
                 # Compute gradients of f_v with respect to the weights
@@ -123,13 +123,13 @@ class WeightsTraceHessianCalculatorPytorch(TraceHessianCalculatorPytorch):
                     approx = torch.sum(approx, dim=shape_channel_axis)
 
                 # Update node Hessian approximation mean over random iterations
-                ipts_hessian_trace_approx[i] = (j * ipts_hessian_trace_approx[i] + approx) / (j + 1)
+                ipts_hessian_approx_scores[i] = (j * ipts_hessian_approx_scores[i] + approx) / (j + 1)
 
             # If the change to the maximal mean Hessian approximation is insignificant we stop the calculation
             # Note that we do not consider granularity when computing the mean
             if j > MIN_HESSIAN_ITER:
                 if prev_mean_results is not None:
-                    new_mean_res = torch.as_tensor([torch.mean(res) for res in ipts_hessian_trace_approx],
+                    new_mean_res = torch.as_tensor([torch.mean(res) for res in ipts_hessian_approx_scores],
                                                    device=device)
                     relative_delta_per_node = (torch.abs(new_mean_res - prev_mean_results) /
                                                (torch.abs(new_mean_res) + 1e-6))
@@ -137,16 +137,16 @@ class WeightsTraceHessianCalculatorPytorch(TraceHessianCalculatorPytorch):
                     if max_delta < HESSIAN_COMP_TOLERANCE:
                         break
 
-            prev_mean_results = torch.as_tensor([torch.mean(res) for res in ipts_hessian_trace_approx], device=device)
+            prev_mean_results = torch.as_tensor([torch.mean(res) for res in ipts_hessian_approx_scores], device=device)
 
         # Make sure all final shape are tensors and not scalar
-        if self.hessian_request.granularity == HessianInfoGranularity.PER_TENSOR:
-            ipts_hessian_trace_approx = [final_approx.reshape(1) for final_approx in ipts_hessian_trace_approx]
+        if self.hessian_request.granularity == HessianScoresGranularity.PER_TENSOR:
+            ipts_hessian_approx_scores = [final_approx.reshape(1) for final_approx in ipts_hessian_approx_scores]
 
         # Add a batch axis to the Hessian approximation tensor (to align with the expected returned shape).
         # We assume per-image computation, so the batch axis size is 1.
         final_approx = [r_final_approx[np.newaxis, ...].detach().cpu().numpy()
-                        for r_final_approx in ipts_hessian_trace_approx]
+                        for r_final_approx in ipts_hessian_approx_scores]
 
         return final_approx
 

model_compression_toolkit/core/pytorch/pytorch_implementation.py

@@ -30,7 +30,7 @@ from model_compression_toolkit.core import QuantizationConfig, FrameworkInfo, Co
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common import Graph, BaseNode
 from model_compression_toolkit.core.common.framework_implementation import FrameworkImplementation
-from model_compression_toolkit.core.common.hessian import TraceHessianRequest, HessianMode, HessianInfoService
+from model_compression_toolkit.core.common.hessian import HessianScoresRequest, HessianMode, HessianInfoService
 from model_compression_toolkit.core.common.mixed_precision.sensitivity_evaluation import SensitivityEvaluation
 from model_compression_toolkit.core.common.mixed_precision.set_layer_to_bitwidth import set_layer_to_bitwidth
 from model_compression_toolkit.core.common.model_builder_mode import ModelBuilderMode
@@ -76,10 +76,10 @@ from model_compression_toolkit.core.pytorch.graph_substitutions.substitutions.we
     WeightsActivationSplit
 from model_compression_toolkit.core.pytorch.graph_substitutions.substitutions.concat_threshold_update import \
     ConcatThresholdUpdate
-from model_compression_toolkit.core.pytorch.hessian.activation_trace_hessian_calculator_pytorch import \
-    ActivationTraceHessianCalculatorPytorch
-from model_compression_toolkit.core.pytorch.hessian.weights_trace_hessian_calculator_pytorch import \
-    WeightsTraceHessianCalculatorPytorch
+from model_compression_toolkit.core.pytorch.hessian.activation_hessian_scores_calculator_pytorch import \
+    ActivationHessianScoresCalculatorPytorch
+from model_compression_toolkit.core.pytorch.hessian.weights_hessian_scores_calculator_pytorch import \
+    WeightsHessianScoresCalculatorPytorch
 from model_compression_toolkit.core.pytorch.mixed_precision.configurable_activation_quantizer import \
     ConfigurableActivationQuantizer
 from model_compression_toolkit.core.pytorch.mixed_precision.configurable_weights_quantizer import \
@@ -353,7 +353,7 @@ class PytorchImplementation(FrameworkImplementation):
             representative_data_gen: Dataset to use for retrieving images for the models inputs.
             fw_info: FrameworkInfo object with information about the specific framework's model.
             disable_activation_for_metric: Whether to disable activation quantization when computing the MP metric.
-            hessian_info_service: HessianInfoService to fetch approximations of the hessian traces for the float model.
+            hessian_info_service: HessianScoresService to fetch approximations of the hessian scores for the float model.
 
         Returns:
             A SensitivityEvaluation object.
@@ -515,34 +515,34 @@ class PytorchImplementation(FrameworkImplementation):
 
         return model(*inputs)
 
-    def get_trace_hessian_calculator(self,
-                                     graph: Graph,
-                                     input_images: List[Any],
-                                     trace_hessian_request: TraceHessianRequest,
-                                     num_iterations_for_approximation: int = HESSIAN_NUM_ITERATIONS):
+    def get_hessian_scores_calculator(self,
+                                      graph: Graph,
+                                      input_images: List[Any],
+                                      hessian_scores_request: HessianScoresRequest,
+                                      num_iterations_for_approximation: int = HESSIAN_NUM_ITERATIONS):
         """
-        Get Pytorch trace hessian approximations calculator based on the trace hessian request.
+        Get Pytorch hessian scores calculator based on the hessian scores request.
         Args:
             input_images: Images to use for computation.
             graph: Float graph to compute the approximation of its different nodes.
-            trace_hessian_request: TraceHessianRequest to search for the desired calculator.
-            num_iterations_for_approximation: Number of iterations to use when approximating the Hessian trace.
-
-        Returns: TraceHessianCalculatorPytorch to use for the trace hessian approximation computation for this request.
-
-        """
-        if trace_hessian_request.mode == HessianMode.ACTIVATION:
-            return ActivationTraceHessianCalculatorPytorch(graph=graph,
-                                                           trace_hessian_request=trace_hessian_request,
-                                                           input_images=input_images,
-                                                           fw_impl=self,
-                                                           num_iterations_for_approximation=num_iterations_for_approximation)
-        elif trace_hessian_request.mode == HessianMode.WEIGHTS:
-            return WeightsTraceHessianCalculatorPytorch(graph=graph,
-                                                        trace_hessian_request=trace_hessian_request,
-                                                        input_images=input_images,
-                                                        fw_impl=self,
-                                                        num_iterations_for_approximation=num_iterations_for_approximation)
+            hessian_scores_request: HessianScoresRequest to search for the desired calculator.
+            num_iterations_for_approximation: Number of iterations to use when approximating the Hessian scores.
+
+        Returns: HessianScoresCalculatorPytorch to use for the hessian approximation scores computation for this request.
+
+        """
+        if hessian_scores_request.mode == HessianMode.ACTIVATION:
+            return ActivationHessianScoresCalculatorPytorch(graph=graph,
+                                                            hessian_scores_request=hessian_scores_request,
+                                                            input_images=input_images,
+                                                            fw_impl=self,
+                                                            num_iterations_for_approximation=num_iterations_for_approximation)
+        elif hessian_scores_request.mode == HessianMode.WEIGHTS:
+            return WeightsHessianScoresCalculatorPytorch(graph=graph,
+                                                         hessian_scores_request=hessian_scores_request,
+                                                         input_images=input_images,
+                                                         fw_impl=self,
+                                                         num_iterations_for_approximation=num_iterations_for_approximation)
 
     def get_inferable_quantizers(self, node: BaseNode):
         """

model_compression_toolkit/core/quantization_prep_runner.py

@@ -40,7 +40,7 @@ def quantization_preparation_runner(graph: Graph,
                                     fw_info: FrameworkInfo,
                                     fw_impl: FrameworkImplementation,
                                     tb_w: TensorboardWriter = None,
-                                    hessian_info_service: HessianInfoService = None,) -> Graph:
+                                    hessian_info_service: HessianInfoService = None, ) -> Graph:
     """
     Prepares a trained model for post-training quantization.
     First, the model graph is optimized using several transformations (e.g. folding BatchNormalization to preceding layers).