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

model_compression_toolkit/core/common/hessian/{trace_hessian_calculator.py → hessian_scores_calculator.py}

@@ -18,16 +18,16 @@ from typing import List, Any
 
 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 import HessianScoresRequest
 from model_compression_toolkit.logger import Logger
 
 
-class TraceHessianCalculator(ABC):
+class HessianScoresCalculator(ABC):
     """
-    Abstract base class for computing an approximation of the trace of the Hessian.
+    Abstract base class for computing scores based on the Hessian matrix approximation.
 
     This class provides a structure for implementing different methods to compute
-    the trace of the Hessian approximation based on the provided configuration,
+    scores based on Hessian-approximation according to the provided configuration,
     input images, and other parameters.
     """
 
@@ -35,15 +35,15 @@ class TraceHessianCalculator(ABC):
                  graph: Graph,
                  input_images: List[Any],
                  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-approximation scores computation.
+            hessian_scores_request: Configuration request for which to compute the Hessian-based approximation.
+            num_iterations_for_approximation: Number of iterations to use when approximating the Hessian-approximation scores.
 
         """
         self.graph = graph
@@ -60,15 +60,15 @@ class TraceHessianCalculator(ABC):
             Logger.critical(f"The graph requires {len(graph.get_inputs())} inputs, but the provided representative dataset contains {len(self.input_images)} inputs.")
 
         self.fw_impl = fw_impl
-        self.hessian_request = trace_hessian_request
+        self.hessian_request = hessian_scores_request
 
     @abstractmethod
     def compute(self) -> List[float]:
         """
-        Abstract method to compute the approximation of the trace of the Hessian.
+        Abstract method to compute the scores based on the Hessian-approximation matrix.
 
         This method should be implemented by subclasses to provide the specific
-        computation method for the trace Hessian approximation.
+        computation method for the Hessian-approximation scores.
         """
         raise NotImplemented(f'{self.__class__.__name__} have to implement compute method.')  # pragma: no cover
 

model_compression_toolkit/core/common/hessian/{trace_hessian_request.py → hessian_scores_request.py}

@@ -28,11 +28,11 @@ class HessianMode(Enum):
     ACTIVATION = 1     # Hessian approximation based on activations
 
 
-class HessianInfoGranularity(Enum):
+class HessianScoresGranularity(Enum):
     """
-    Enum representing the granularity level for Hessian information computation.
+    Enum representing the granularity level for Hessian scores computation.
 
-    This determines the number the Hessian approximations is computed for some node.
+    This determines the number the Hessian scores is computed for some node.
     Note: This is not the actual Hessian but an approximation.
     """
     PER_ELEMENT = 0
@@ -40,25 +40,25 @@ class HessianInfoGranularity(Enum):
     PER_TENSOR = 2
 
 
-class TraceHessianRequest:
+class HessianScoresRequest:
     """
-    Request configuration for the trace of the Hessian approximation.
+    Request configuration for the Hessian-approximation scores.
 
-    This class defines the parameters for the approximation of the trace of the Hessian matrix.
+    This class defines the parameters for the scores based on the Hessian matrix approximation.
     It specifies the mode (weights/activations), granularity (element/channel/tensor), and the target node.
-    Note: This does not compute the actual Hessian's trace but approximates it.
+
+    Note: This does not compute scores using the actual Hessian matrix but an approximation.
     """
 
     def __init__(self,
                  mode: HessianMode,
-                 granularity: HessianInfoGranularity,
-                 target_nodes: List,
-                 ):
+                 granularity: HessianScoresGranularity,
+                 target_nodes: List):
         """
         Attributes:
-            mode (HessianMode): Mode of Hessian's trace approximation (w.r.t weights or activations).
-            granularity (HessianInfoGranularity): Granularity level for the approximation.
-            target_nodes (List[BaseNode]): The node in the float graph for which the Hessian's trace approximation is targeted.
+            mode (HessianMode): Mode of Hessian-approximation score (w.r.t weights or activations).
+            granularity (HessianScoresGranularity): Granularity level for the approximation.
+            target_nodes (List[BaseNode]): The node in the float graph for which the Hessian's approximation scores is targeted.
         """
 
         self.mode = mode  # w.r.t activations or weights
@@ -66,9 +66,9 @@ class TraceHessianRequest:
         self.target_nodes = target_nodes
 
     def __eq__(self, other):
-        # Checks if the other object is an instance of TraceHessianRequest
+        # Checks if the other object is an instance of HessianScoresRequest
         # and then checks if all attributes are equal.
-        return isinstance(other, TraceHessianRequest) and \
+        return isinstance(other, HessianScoresRequest) and \
                self.mode == other.mode and \
                self.granularity == other.granularity and \
                self.target_nodes == other.target_nodes

model_compression_toolkit/core/common/mixed_precision/mixed_precision_search_facade.py

@@ -51,7 +51,7 @@ def search_bit_width(graph_to_search_cfg: Graph,
                      mp_config: MixedPrecisionQuantizationConfig,
                      representative_data_gen: Callable,
                      search_method: BitWidthSearchMethod = BitWidthSearchMethod.INTEGER_PROGRAMMING,
-                     hessian_info_service: HessianInfoService=None) -> List[int]:
+                     hessian_info_service: HessianInfoService = None) -> List[int]:
     """
     Search for an MP configuration for a given graph. Given a search_method method (by default, it's linear
     programming), we use the sensitivity_evaluator object that provides a function to compute an
@@ -68,7 +68,7 @@ def search_bit_width(graph_to_search_cfg: Graph,
         mp_config: Mixed-precision quantization configuration.
         representative_data_gen: Dataset to use for retrieving images for the models inputs.
         search_method: BitWidthSearchMethod to define which searching method to use.
-        hessian_info_service: HessianInfoService to fetch Hessian traces approximations.
+        hessian_info_service: HessianInfoService to fetch Hessian-approximation information.
 
     Returns:
         A MP configuration for the graph (list of integers, where the index in the list, is the node's

model_compression_toolkit/core/common/mixed_precision/sensitivity_evaluation.py

@@ -24,8 +24,8 @@ from model_compression_toolkit.core.common.graph.functional_node import Function
 from model_compression_toolkit.core.common.similarity_analyzer import compute_kl_divergence
 from model_compression_toolkit.core.common.model_builder_mode import ModelBuilderMode
 from model_compression_toolkit.logger import Logger
-from model_compression_toolkit.core.common.hessian import TraceHessianRequest, HessianMode, \
-    HessianInfoGranularity, HessianInfoService
+from model_compression_toolkit.core.common.hessian import HessianScoresRequest, HessianMode, \
+    HessianScoresGranularity, HessianInfoService
 
 
 class SensitivityEvaluation:
@@ -65,7 +65,7 @@ class SensitivityEvaluation:
             set_layer_to_bitwidth: A fw-dependent function that allows to configure a configurable MP model
                     with a specific bit-width configuration.
             disable_activation_for_metric: Whether to disable activation quantization when computing the MP metric.
-            hessian_info_service: HessianInfoService to fetch Hessian traces approximations.
+            hessian_info_service: HessianInfoService to fetch Hessian approximation information.
 
         """
         self.graph = graph
@@ -237,14 +237,14 @@ class SensitivityEvaluation:
          to be used for the distance metric weighted average computation.
 
         """
-        # Create a request for trace Hessian approximation with specific configurations
+        # Create a request for Hessian approximation scores with specific configurations
         # (here we use per-tensor approximation of the Hessian's trace w.r.t the node's activations)
-        trace_hessian_request = TraceHessianRequest(mode=HessianMode.ACTIVATION,
-                                                    granularity=HessianInfoGranularity.PER_TENSOR,
+        hessian_info_request = HessianScoresRequest(mode=HessianMode.ACTIVATION,
+                                                    granularity=HessianScoresGranularity.PER_TENSOR,
                                                     target_nodes=self.interest_points)
 
-        # Fetch the trace Hessian approximations for the current interest point
-        nodes_approximations = self.hessian_info_service.fetch_hessian(trace_hessian_request=trace_hessian_request,
+        # Fetch the Hessian approximation scores for the current interest point
+        nodes_approximations = self.hessian_info_service.fetch_hessian(hessian_scores_request=hessian_info_request,
                                                                        required_size=self.quant_config.num_of_images,
                                                                        batch_size=self.quant_config.hessian_batch_size)
 

model_compression_toolkit/core/common/pruning/importance_metrics/lfh_importance_metric.py

@@ -19,8 +19,8 @@ from typing import Callable, List, Dict, Tuple
 
 from model_compression_toolkit.core.common import Graph, BaseNode
 from model_compression_toolkit.core.common.framework_info import FrameworkInfo
-from model_compression_toolkit.core.common.hessian import HessianInfoService, HessianMode, HessianInfoGranularity, \
-    TraceHessianRequest
+from model_compression_toolkit.core.common.hessian import HessianInfoService, HessianMode, HessianScoresGranularity, \
+    HessianScoresRequest
 from model_compression_toolkit.core.common.pruning.channels_grouping import ChannelGrouping
 from model_compression_toolkit.core.common.pruning.importance_metrics.base_importance_metric import BaseImportanceMetric
 from model_compression_toolkit.core.common.pruning.pruning_config import PruningConfig
@@ -127,9 +127,9 @@ class LFHImportanceMetric(BaseImportanceMetric):
         # Fetch and process Hessian scores for output channels of entry nodes.
         nodes_scores = []
         for node in entry_nodes:
-            _request = TraceHessianRequest(mode=HessianMode.WEIGHTS,
-                                           granularity=HessianInfoGranularity.PER_OUTPUT_CHANNEL,
-                                           target_nodes=[node])
+            _request = HessianScoresRequest(mode=HessianMode.WEIGHTS,
+                                            granularity=HessianScoresGranularity.PER_OUTPUT_CHANNEL,
+                                            target_nodes=[node])
             _scores_for_node = hessian_info_service.fetch_hessian(_request,
                                                                   required_size=self.pruning_config.num_score_approximations)
             nodes_scores.append(_scores_for_node)

model_compression_toolkit/core/common/quantization/quantization_params_generation/error_functions.py

@@ -16,7 +16,7 @@ from copy import deepcopy
 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, \
+from model_compression_toolkit.core.common.hessian import HessianScoresRequest, HessianMode, HessianScoresGranularity, \
     HessianInfoService
 from model_compression_toolkit.core.common.similarity_analyzer import compute_mse, compute_mae, compute_lp_norm
 from model_compression_toolkit.logger import Logger
@@ -389,9 +389,9 @@ def _compute_hessian_for_hmse(node,
     Returns: A list with computed Hessian-based scores tensors for the given node.
 
     """
-    _request = TraceHessianRequest(mode=HessianMode.WEIGHTS,
-                                   granularity=HessianInfoGranularity.PER_ELEMENT,
-                                   target_nodes=[node])
+    _request = HessianScoresRequest(mode=HessianMode.WEIGHTS,
+                                    granularity=HessianScoresGranularity.PER_ELEMENT,
+                                    target_nodes=[node])
     _scores_for_node = hessian_info_service.fetch_hessian(_request,
                                                           required_size=num_hessian_samples)
 

model_compression_toolkit/core/common/quantization/quantization_params_generation/qparams_computation.py

@@ -20,8 +20,8 @@ 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, TraceHessianRequest, HessianMode, \
-    HessianInfoGranularity
+from model_compression_toolkit.core.common.hessian import HessianInfoService, HessianScoresRequest, HessianMode, \
+    HessianScoresGranularity
 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 \
@@ -89,9 +89,9 @@ def calculate_quantization_params(graph: Graph,
     # The Hessian scores are computed and stored in the hessian_info_service object.
     nodes_for_hmse = _collect_nodes_for_hmse(nodes_list, graph)
     if len(nodes_for_hmse) > 0:
-        hessian_info_service.fetch_hessian(TraceHessianRequest(mode=HessianMode.WEIGHTS,
-                                                               granularity=HessianInfoGranularity.PER_ELEMENT,
-                                                               target_nodes=nodes_for_hmse),
+        hessian_info_service.fetch_hessian(HessianScoresRequest(mode=HessianMode.WEIGHTS,
+                                                                granularity=HessianScoresGranularity.PER_ELEMENT,
+                                                                target_nodes=nodes_for_hmse),
                                            required_size=num_hessian_samples,
                                            batch_size=1)
 

model_compression_toolkit/core/keras/hessian/{activation_trace_hessian_calculator_keras.py → activation_hessian_scores_calculator_keras.py}

@@ -22,45 +22,45 @@ 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.keras.back2framework.float_model_builder import FloatKerasModelBuilder
-from model_compression_toolkit.core.keras.hessian.trace_hessian_calculator_keras import TraceHessianCalculatorKeras
+from model_compression_toolkit.core.keras.hessian.hessian_scores_calculator_keras import HessianScoresCalculatorKeras
 from model_compression_toolkit.logger import Logger
 
 
-class ActivationTraceHessianCalculatorKeras(TraceHessianCalculatorKeras):
+class ActivationHessianScoresCalculatorKeras(HessianScoresCalculatorKeras):
     """
-    Keras implementation of the Trace Hessian Calculator for activations.
+    Keras implementation of the Hessian-approximation scores Calculator for activations.
     """
     def __init__(self,
                  graph: Graph,
                  input_images: List[tf.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(ActivationTraceHessianCalculatorKeras, 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(ActivationHessianScoresCalculatorKeras, 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 Hessian-approximation based scores 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 Hessian-approximation 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()]
 
             if len([n for n in self.hessian_request.target_nodes if n in model_output_nodes]) > 0:
@@ -98,9 +98,9 @@ class ActivationTraceHessianCalculatorKeras(TraceHessianCalculatorKeras):
                 # Unfold and concatenate all outputs to form a single tensor
                 output = self._concat_tensors(output_tensors)
 
-                # List to store the approximated trace of the Hessian for each interest point
-                ipts_hessian_trace_approx = [tf.Variable([0.0], dtype=tf.float32, trainable=True)
-                                             for _ in range(len(target_activation_tensors))]
+                # List to store the Hessian-approximation scores for each interest point
+                ipts_hessian_approximations = [tf.Variable([0.0], dtype=tf.float32, trainable=True)
+                                               for _ in range(len(target_activation_tensors))]
 
                 # Loop through each interest point activation tensor
                 prev_mean_results = None
@@ -120,29 +120,29 @@ class ActivationTraceHessianCalculatorKeras(TraceHessianCalculatorKeras):
                                 continue  # pragma: no cover
 
                             # Mean over all dims but the batch (CXHXW for conv)
-                            hessian_trace_approx = tf.reduce_sum(hess_v ** 2.0,
-                                                                 axis=tuple(d for d in range(1, len(hess_v.shape))))
+                            hessian_approx = tf.reduce_sum(hess_v ** 2.0,
+                                                           axis=tuple(d for d in range(1, len(hess_v.shape))))
 
                             # Free gradients
                             del hess_v
 
                             # 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_approximations[i] = (j * ipts_hessian_approximations[i] + hessian_approx) / (j + 1)
 
                     # If the change to the mean approximation is insignificant (to all outputs)
                     # we stop the calculation.
                     if j > MIN_HESSIAN_ITER:
                         if prev_mean_results is not None:
-                            new_mean_res = tf.reduce_mean(tf.stack(ipts_hessian_trace_approx), axis=1)
+                            new_mean_res = tf.reduce_mean(tf.stack(ipts_hessian_approximations), axis=1)
                             relative_delta_per_node = (tf.abs(new_mean_res - prev_mean_results) /
                                                        (tf.abs(new_mean_res) + 1e-6))
                             max_delta = tf.reduce_max(relative_delta_per_node)
                             if max_delta < HESSIAN_COMP_TOLERANCE:
                                 break
-                    prev_mean_results = tf.reduce_mean(tf.stack(ipts_hessian_trace_approx), axis=1)
+                    prev_mean_results = tf.reduce_mean(tf.stack(ipts_hessian_approximations), axis=1)
 
                 # Convert results to list of numpy arrays
-                hessian_results = [h.numpy() for h in ipts_hessian_trace_approx]
+                hessian_results = [h.numpy() for h in ipts_hessian_approximations]
                 # 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.
@@ -152,4 +152,4 @@ class ActivationTraceHessianCalculatorKeras(TraceHessianCalculatorKeras):
 
         else:  # pragma: no cover
             Logger.critical(f"{self.hessian_request.granularity} "
-                            f"is not supported for Keras activation hessian\'s trace approximation calculator.")
+                            f"is not supported for Keras activation hessian\'s approximation scores calculator.")

model_compression_toolkit/core/keras/hessian/{trace_hessian_calculator_keras.py → hessian_scores_calculator_keras.py}

@@ -13,7 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 
-from model_compression_toolkit.core.common.hessian.trace_hessian_calculator import TraceHessianCalculator
+from model_compression_toolkit.core.common.hessian.hessian_scores_calculator import HessianScoresCalculator
 
 from typing import List, Tuple, Dict, Any, Union
 
@@ -23,38 +23,38 @@ from model_compression_toolkit.constants import HESSIAN_NUM_ITERATIONS
 from model_compression_toolkit.core.common.graph.edge import EDGE_SINK_INDEX
 from model_compression_toolkit.core.common import Graph, BaseNode
 from model_compression_toolkit.core.common.graph.functional_node import FunctionalNode
-from model_compression_toolkit.core.common.hessian import TraceHessianRequest
+from model_compression_toolkit.core.common.hessian import HessianScoresRequest
 from model_compression_toolkit.core.keras.back2framework.instance_builder import OperationHandler
 from tensorflow.python.util.object_identity import Reference as TFReference
 
 from model_compression_toolkit.logger import Logger
 
 
-class TraceHessianCalculatorKeras(TraceHessianCalculator):
+class HessianScoresCalculatorKeras(HessianScoresCalculator):
     """
-    Keras-specific implementation of the Trace Hessian approximation Calculator.
-    This class serves as a base for other Keras-specific trace Hessian approximation calculators.
+    Keras-specific implementation of the Hessian approximation scores Calculator.
+    This class serves as a base for other Keras-specific Hessian approximation scores calculators.
     """
     def __init__(self,
                  graph: Graph,
                  input_images: List[tf.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-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-based scores.
         """
-        super(TraceHessianCalculatorKeras, 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(HessianScoresCalculatorKeras, 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[tf.Tensor, List[tf.Tensor]]) -> tf.Tensor:
         """

model_compression_toolkit/core/keras/hessian/{weights_trace_hessian_calculator_keras.py → weights_hessian_scores_calculator_keras.py}

@@ -20,39 +20,39 @@ from typing import List
 
 from model_compression_toolkit.constants import HESSIAN_NUM_ITERATIONS, MIN_HESSIAN_ITER, HESSIAN_COMP_TOLERANCE
 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.keras.back2framework.float_model_builder import FloatKerasModelBuilder
 from model_compression_toolkit.core.keras.default_framework_info import DEFAULT_KERAS_INFO
-from model_compression_toolkit.core.keras.hessian.trace_hessian_calculator_keras import TraceHessianCalculatorKeras
+from model_compression_toolkit.core.keras.hessian.hessian_scores_calculator_keras import HessianScoresCalculatorKeras
 from model_compression_toolkit.logger import Logger
 
 
-class WeightsTraceHessianCalculatorKeras(TraceHessianCalculatorKeras):
+class WeightsHessianScoresCalculatorKeras(HessianScoresCalculatorKeras):
     """
-    Keras-specific implementation of the Trace Hessian approximation computation w.r.t a node's weights.
+    Keras-specific implementation of the Hessian-approximation scores computation w.r.t a node's weights.
     """
 
     def __init__(self,
                  graph: Graph,
                  input_images: List[tf.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(WeightsTraceHessianCalculatorKeras, 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(WeightsHessianScoresCalculatorKeras, 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]:
         """
@@ -83,7 +83,7 @@ class WeightsTraceHessianCalculatorKeras(TraceHessianCalculatorKeras):
             # Combine outputs if the model returns multiple output tensors
             output = self._concat_tensors(outputs)
 
-            ipts_hessian_trace_approx = [tf.Variable([0.0], dtype=tf.float32, trainable=True)
+            ipts_hessian_scores_approx = [tf.Variable([0.0], dtype=tf.float32, trainable=True)
                                          for _ in range(len(self.hessian_request.target_nodes))]
 
             prev_mean_results = None
@@ -134,7 +134,7 @@ class WeightsTraceHessianCalculatorKeras(TraceHessianCalculatorKeras):
                         approx = tf.reduce_sum(tf.pow(gradients, 2.0), axis=1)
 
                         # Update node Hessian approximation mean over random iterations
-                        ipts_hessian_trace_approx[i] = (j * ipts_hessian_trace_approx[i] + approx) / (j + 1)
+                        ipts_hessian_scores_approx[i] = (j * ipts_hessian_scores_approx[i] + approx) / (j + 1)
 
                         # Free gradients
                         del gradients
@@ -145,33 +145,33 @@ class WeightsTraceHessianCalculatorKeras(TraceHessianCalculatorKeras):
                     if j > MIN_HESSIAN_ITER:
                         if prev_mean_results is not None:
                             new_mean_res = \
-                                tf.convert_to_tensor([tf.reduce_mean(res) for res in ipts_hessian_trace_approx])
+                                tf.convert_to_tensor([tf.reduce_mean(res) for res in ipts_hessian_scores_approx])
                             relative_delta_per_node = (tf.abs(new_mean_res - prev_mean_results) /
                                                        (tf.abs(new_mean_res) + 1e-6))
                             max_delta = tf.reduce_max(relative_delta_per_node)
                             if max_delta < HESSIAN_COMP_TOLERANCE:
                                 break
 
-                prev_mean_results = tf.convert_to_tensor([tf.reduce_mean(res) for res in ipts_hessian_trace_approx])
+                prev_mean_results = tf.convert_to_tensor([tf.reduce_mean(res) for res in ipts_hessian_scores_approx])
 
         # Free gradient tape
         del tape
 
-        if self.hessian_request.granularity == HessianInfoGranularity.PER_TENSOR:
-            for final_approx in ipts_hessian_trace_approx:
+        if self.hessian_request.granularity == HessianScoresGranularity.PER_TENSOR:
+            for final_approx in ipts_hessian_scores_approx:
                 if final_approx.shape != (1,):  # pragma: no cover
                     Logger.critical(f"For HessianInfoGranularity.PER_TENSOR, the expected score shape is (1,), "
                                     f"but found {final_approx.shape}.")
-        elif self.hessian_request.granularity == HessianInfoGranularity.PER_ELEMENT:
+        elif self.hessian_request.granularity == HessianScoresGranularity.PER_ELEMENT:
             # Reshaping the scores to the original weight shape
-            ipts_hessian_trace_approx = \
+            ipts_hessian_scores_approx = \
                 [tf.reshape(final_approx, s) for final_approx, s in
-                 zip(ipts_hessian_trace_approx, tensors_original_shape)]
+                 zip(ipts_hessian_scores_approx, tensors_original_shape)]
 
         # 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, ...].numpy()
-                        for r_final_approx in ipts_hessian_trace_approx]
+                        for r_final_approx in ipts_hessian_scores_approx]
 
         return final_approx
 
@@ -193,7 +193,7 @@ class WeightsTraceHessianCalculatorKeras(TraceHessianCalculatorKeras):
             tf.Tensor: Reshaped gradient tensor based on the granularity.
         """
         # Reshape the gradients based on the granularity (whole tensor, per channel, or per element)
-        if self.hessian_request.granularity != HessianInfoGranularity.PER_OUTPUT_CHANNEL:
+        if self.hessian_request.granularity != HessianScoresGranularity.PER_OUTPUT_CHANNEL:
             gradients = tf.reshape(gradients, [num_of_scores, -1])
         else:
             # Slice the gradients, vectorize them and stack them along the first axis.
@@ -216,11 +216,11 @@ class WeightsTraceHessianCalculatorKeras(TraceHessianCalculatorKeras):
         Returns:
             int: The number of scores.
         """
-        if self.hessian_request.granularity == HessianInfoGranularity.PER_TENSOR:
+        if self.hessian_request.granularity == HessianScoresGranularity.PER_TENSOR:
             return 1
-        elif self.hessian_request.granularity == HessianInfoGranularity.PER_OUTPUT_CHANNEL:
+        elif self.hessian_request.granularity == HessianScoresGranularity.PER_OUTPUT_CHANNEL:
             return weight_tensor.shape[output_channel_axis]
-        elif self.hessian_request.granularity == HessianInfoGranularity.PER_ELEMENT:
+        elif self.hessian_request.granularity == HessianScoresGranularity.PER_ELEMENT:
             return tf.size(weight_tensor).numpy()
         else:  # pragma: no cover
             Logger.critical(f"Unexpected granularity encountered: {self.hessian_request.granularity}.")