mct-nightly 2.3.0.20250224.520__py3-none-any.whl → 2.3.0.20250225.512__py3-none-any.whl

model_compression_toolkit/core/common/model_collector.py

@@ -15,13 +15,15 @@
 
 
 import numpy as np
-from typing import List
+from typing import List, Union, Tuple, Optional
 
 from networkx.algorithms.dag import topological_sort
-from model_compression_toolkit.core import FrameworkInfo
+from model_compression_toolkit.core import FrameworkInfo, QuantizationErrorMethod
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common.framework_implementation import FrameworkImplementation
 from model_compression_toolkit.core.common.graph.base_graph import Graph
+from model_compression_toolkit.core.common.hessian import HessianInfoService, HessianScoresGranularity, HessianMode, \
+    HessianScoresRequest
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.core.common.model_builder_mode import ModelBuilderMode
 from model_compression_toolkit.core.common.collectors.statistics_collector import BaseStatsCollector
@@ -66,12 +68,67 @@ def create_tensor2node(graph: common.Graph,
 
     """
     current_sc = graph.get_out_stats_collector(node)
-    is_list_nostat_collectors = isinstance(current_sc, list) and len([sc for sc in current_sc if not isinstance(sc, common.NoStatsCollector)]) == 0
+    is_list_nostat_collectors = isinstance(current_sc, list) and len(
+        [sc for sc in current_sc if not isinstance(sc, common.NoStatsCollector)]) == 0
     if isinstance(current_sc, common.NoStatsCollector) or current_sc is None or is_list_nostat_collectors:
         stats_collector = common.StatsCollector(fw_info.out_channel_axis_mapping.get(node.type))
         graph.set_out_stats_collector_to_node(node, stats_collector)
 
 
+def ensure_matching_data_lengths(
+    stats_collector: Union[List[BaseStatsCollector], Tuple[BaseStatsCollector, ...]],
+    tensor_data: Union[List, Tuple],
+    hessian_data: Union[List, Tuple]
+):
+    """
+    Ensures that the lengths of `tensor_data`, `hessian_data`, and `stats_collector` are matching.
+    If the types or lengths do not match, a critical error is logged.
+
+    Args:
+        stats_collector: A list or tuple of statistics collectors.
+        tensor_data: A list or tuple of tensors corresponding to the statistics collectors.
+        hessian_data: A list or tuple of Hessian tensors corresponding to the statistics collectors.
+
+    Raises:
+        Logs a critical error and halts execution if there is a type mismatch or
+        if the lengths of the inputs do not match.
+    """
+
+    if not isinstance(tensor_data, (list, tuple)):
+        Logger.critical(
+            f"'tensor_data' is of type {type(tensor_data)}, but must be of the same type as 'stats_collector' ({type(stats_collector)})."
+        )  # pragma: no cover
+
+    if len(stats_collector) != len(tensor_data):
+        Logger.critical(
+            "'tensor_data' and 'stats_collector' must have matching lengths."
+        )  # pragma: no cover
+
+    if not isinstance(hessian_data, (list, tuple)):
+        Logger.critical(
+            f"'hessian_data' is of type {type(hessian_data)}, but must be of the same type as 'stats_collector' ({type(stats_collector)})."
+        )  # pragma: no cover
+
+    if len(stats_collector) != len(hessian_data):
+        Logger.critical(
+            "'hessian_data' and 'stats_collector' must have matching lengths."
+        )  # pragma: no cover
+
+
+def convert_to_numpy_and_abs(tensor: Optional[np.ndarray], fw_impl: FrameworkImplementation) -> Optional[np.ndarray]:
+    """
+    Converts a tensor to a NumPy array and applies the absolute value operation if the tensor is not None.
+
+    Args:
+        tensor: Input tensor to be converted to a NumPy array.
+        fw_impl: Framework implementation that provides the 'to_numpy' method for tensor conversion.
+
+    Returns:
+        A NumPy array of the input tensor with absolute values applied. If the input tensor is None, returns None.
+    """
+    return tensor if tensor is None else np.abs(fw_impl.to_numpy(tensor))
+
+
 class ModelCollector:
     """
     Build a model from a graph for statistics collection purposes.
@@ -83,6 +140,7 @@ class ModelCollector:
     def __init__(self, graph: Graph,
                  fw_impl: FrameworkImplementation,
                  fw_info: FrameworkInfo,
+                 hessian_info_service: HessianInfoService = None,
                  qc: common.QuantizationConfig = common.DEFAULTCONFIG):
         """
         Build a model from a graph per framework for statistics collection.
@@ -96,14 +154,18 @@ class ModelCollector:
 
         self.fw_impl = fw_impl
         self.fw_info = fw_info
+        self.hessian_service = hessian_info_service
+        self.qc = qc
+        self.model_outputs = [out.node for out in graph.get_outputs()]
 
-        # Assign statisitcs collectors to nodes
+        # Assign statistics collectors to nodes
         for n in graph.get_topo_sorted_nodes():
             sc = create_stats_collector_for_node(n, fw_info=fw_info)  # Get static collector for the node
             # If we use bias correction, and the node has kernel weights to quantize, we need to make sure
             # its previous nodes' tensors are consistent with this node.
             kernel_attr = fw_info.get_kernel_op_attributes(n.type)[0]
-            if qc.weights_bias_correction and kernel_attr is not None and n.is_weights_quantization_enabled(kernel_attr):
+            if qc.weights_bias_correction and kernel_attr is not None and n.is_weights_quantization_enabled(
+                    kernel_attr):
                 for ie in graph.incoming_edges(n):
                     input_node = ie.source_node
                     create_tensor2node(graph,
@@ -112,7 +174,6 @@ class ModelCollector:
             if sc is not None:
                 graph.set_out_stats_collector_to_node(n, sc)
 
-
         outputs_nodes = []  # List of graph nodes, the model should output their outputs.
         self.stats_containers_list = []  # List of output statistics containers of nodes ordered
         # the same as outputs_nodes so statistics of outputs can be gathered for the correct statistics container.
@@ -135,11 +196,19 @@ class ModelCollector:
                     outputs_nodes.append(n)
                     self.stats_containers_list.append(out_stats_container)
 
+        self.intermediate_output_tensors = [n for n in outputs_nodes if n not in self.model_outputs]
+
+        # Append nodes from graph.get_outputs() that are not already in outputs_nodes for Hessian
+        # calculation for output nodes that don't collect statistics such as "permute", "transpose" etc.
+        # TODO: Add integration test for this case
+        append2output = outputs_nodes + [n for n in self.model_outputs if n not in outputs_nodes]
+
+
         # Build a float model and output all layers' outputs
         # (that should be collected) as the model's outputs
         self.model, _ = self.fw_impl.model_builder(graph,
                                                    mode=ModelBuilderMode.FLOAT,
-                                                   append2output=outputs_nodes,
+                                                   append2output=append2output,
                                                    fw_info=self.fw_info)
 
     def infer(self, inputs_list: List[np.ndarray]):
@@ -154,14 +223,43 @@ class ModelCollector:
 
         # TODO: Thinking about delegating collections to framework
         # TODO: migrate datasets to framework datasets
-        tensor_data = self.fw_impl.run_model_inference(self.model, inputs_list)
-        for td, sc in zip(tensor_data, self.stats_containers_list):
-            if isinstance(sc, (list, tuple)):
-                if not isinstance(td, (list, tuple)):
-                    Logger.critical(f"\'tensor_data\' is of type {type(td)} but must be of the same type as \'stats_containers_list\', which is of type {type(sc)}") # pragma: no cover
-                if len(sc) != len(td):
-                    Logger.critical('\'tensor_data\' and \'stats_containers_list\' must have matching lengths') # pragma: no cover
-                for tdi, sci in zip(td, sc):
-                    sci.update_statistics(self.fw_impl.to_numpy(tdi))
+        compute_hessians = self.qc.activation_error_method == QuantizationErrorMethod.HMSE
+
+        # Retrieve intermediate layer activations for statistical analysis.
+        # Enable gradient computation if Hessian calculations are required.
+        activation_tensors = self.fw_impl.run_model_inference(self.model, inputs_list, requires_grad=compute_hessians)
+
+        if compute_hessians:
+            if self.hessian_service is None:
+                Logger.critical(
+                    "Hessian computation is enabled but `hessian_service` is not initialized. "
+                    "Ensure that `hessian_service` is properly set."
+                )  # pragma: no cover
+            request = HessianScoresRequest(
+                mode=HessianMode.ACTIVATION,
+                granularity=HessianScoresGranularity.PER_ELEMENT,
+                target_nodes=self.intermediate_output_tensors,
+                data_loader=None,
+                n_samples=None,
+                compute_from_tensors=True
+            )
+            hessian_tensors = self.hessian_service.fetch_hessian(request=request,
+                                                                 activation_tensors=activation_tensors)
+            hessian_tensors = list(hessian_tensors.values())
+        else:
+            hessian_tensors = []
+
+        # Hessian is not calculated for the output, add "None" as weights for output tenosrs
+        hessian_tensors += [None for _ in range(len(activation_tensors) - len(hessian_tensors))]
+
+        for activation_tensor, hessian_tensor, stats_container in zip(activation_tensors, hessian_tensors, self.stats_containers_list):
+            if isinstance(stats_container, (list, tuple)):
+                if hessian_tensor is None:
+                    hessian_tensor = [None for _ in range(len(activation_tensor))]
+                ensure_matching_data_lengths(activation_tensor, hessian_tensor, stats_container)
+                for activation_tensor_i, hessian_tensor_i, sci in zip(activation_tensor, hessian_tensor, stats_container):
+                    sci.update_statistics(self.fw_impl.to_numpy(activation_tensor_i),
+                                          convert_to_numpy_and_abs(hessian_tensor_i, self.fw_impl))
             else:
-                sc.update_statistics(self.fw_impl.to_numpy(td))
+                stats_container.update_statistics(self.fw_impl.to_numpy(activation_tensor),
+                                                  convert_to_numpy_and_abs(hessian_tensor, self.fw_impl))

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

@@ -514,6 +514,8 @@ def get_threshold_selection_histogram_error_function(quantization_method: Quanti
     quant_method_error_function_mapping = {
         qc.QuantizationErrorMethod.MSE: lambda q_bins, q_count, bins, counts, threshold, _range:
         _mse_error_histogram(q_bins, q_count, bins, counts),
+        qc.QuantizationErrorMethod.HMSE: lambda q_bins, q_count, bins, counts, threshold, _range:
+        _mse_error_histogram(q_bins, q_count, bins, counts), #HMSE need the same functionality as MSE
         qc.QuantizationErrorMethod.MAE: lambda q_bins, q_count, bins, counts, threshold, _range:
         _mae_error_histogram(q_bins, q_count, bins, counts),
         qc.QuantizationErrorMethod.LP: lambda q_bins, q_count, bins, counts, threshold, _range:

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

@@ -13,66 +13,143 @@
 # limitations under the License.
 # ==============================================================================
 import numpy as np
-from typing import Dict, Union
+from typing import Dict, Union, Optional, Tuple
 
 from mct_quantizers import QuantizationMethod
+from model_compression_toolkit.core import QuantizationErrorMethod
 from model_compression_toolkit.target_platform_capabilities.schema.mct_current_schema import Signedness
 from model_compression_toolkit.core.common.collectors.statistics_collector import BaseStatsCollector
 from model_compression_toolkit.core.common.quantization import quantization_params_generation
 from model_compression_toolkit.core.common.node_prior_info import NodePriorInfo
 from model_compression_toolkit.core.common.quantization.node_quantization_config import NodeActivationQuantizationConfig
 
-
-def get_activations_qparams(activation_quant_cfg: NodeActivationQuantizationConfig,
-                            nodes_prior_info: NodePriorInfo,
-                            out_stats_container: BaseStatsCollector) -> Dict[str, Union[np.ndarray, float, bool]]:
+def get_histogram_data(
+    activation_quant_cfg: NodeActivationQuantizationConfig,
+    out_stats_container: BaseStatsCollector
+) -> Tuple[Optional[np.ndarray], Optional[np.ndarray]]:
     """
-    Compute the activations params for a given node in a graph according to a params function.
+    Extract and filter the histogram data from the statistics container.
 
     Args:
-        activation_quant_cfg: node's activation quantization configuration.
-        nodes_prior_info: Prior info collected for the node that is being quantized.
-        out_stats_container: Tensor containing output statistics of the node.
+        activation_quant_cfg: Node's activation quantization configuration.
+        out_stats_container: Statistics container with histogram data.
 
     Returns:
-        The computed activation quantization params.
+        A tuple containing the filtered bins_values and bins_counts.
     """
-
     bins_values, bins_counts = None, None
 
     # If the statistics container collected the histogram, we start by filtering outliers using z threshold
     # filtering, and then computing the threshold based on the filtered histogram.
     if out_stats_container.require_collection():
-        bins_values, bins_counts = out_stats_container.hc.get_histogram()
-        bins_counts = quantization_params_generation.z_score_filter(activation_quant_cfg.z_threshold,
-                                                                    bins_values,
-                                                                    bins_counts)
-    min_value, max_value = out_stats_container.get_min_max_values()
+        if activation_quant_cfg.activation_error_method == QuantizationErrorMethod.HMSE:
+            bins_values, bins_counts = out_stats_container.weighted_hc.get_histogram()
+        else:
+            bins_values, bins_counts = out_stats_container.hc.get_histogram()
+        bins_counts = quantization_params_generation.z_score_filter(
+            activation_quant_cfg.z_threshold,
+            bins_values,
+            bins_counts
+        )
+    return bins_values, bins_counts
+
+def determine_signedness(
+    activation_quant_cfg: NodeActivationQuantizationConfig,
+    nodes_prior_info: NodePriorInfo,
+    min_value: float,
+    bins_values: Optional[np.ndarray],
+    bins_counts: Optional[np.ndarray]
+) -> bool:
+    """
+    Determine if the activations should be considered signed based on the quantization configuration,
+    node prior information, and histogram statistics.
+
+    Args:
+        activation_quant_cfg: Node's activation quantization configuration.
+        nodes_prior_info: Prior info collected for the node that is being quantized.
+        min_value: Minimum value from the statistics container.
+        bins_values: Numpy array of histogram bin values.
+        bins_counts: Numpy array of histogram bin counts.
 
+    Returns:
+        A boolean indicating if the activations are signed.
+    """
     if activation_quant_cfg.signedness in [Signedness.SIGNED, Signedness.UNSIGNED]:
-        signed = activation_quant_cfg.signedness == Signedness.SIGNED
-    elif nodes_prior_info.is_output_bounded():
-        signed = min_value < 0
-    else:
-        signed = np.any(bins_values[:-1][bins_counts > 0] < 0)
+        return activation_quant_cfg.signedness == Signedness.SIGNED
 
+    if nodes_prior_info.is_output_bounded():
+        return min_value < 0
+
+    return np.any(bins_values[:-1][bins_counts > 0] < 0)
+
+
+def update_activation_quantization_params_fn(
+        activation_quant_cfg: NodeActivationQuantizationConfig,
+        nodes_prior_info: NodePriorInfo):
+    """
+    Update the activation quantization parameters function based on the quantization method
+    and whether the node's output is bounded.
+
+    Args:
+        activation_quant_cfg: Node's activation quantization configuration.
+        nodes_prior_info: Prior info collected for the node that is being quantized.
+    """
     if nodes_prior_info.is_output_bounded():
         if activation_quant_cfg.activation_quantization_method == QuantizationMethod.POWER_OF_TWO:
-            activation_quant_cfg.activation_quantization_params_fn = \
+            activation_quant_cfg.set_activation_quantization_params_fn(
                 quantization_params_generation.power_of_two_no_clipping_selection_min_max
+            )
         elif activation_quant_cfg.activation_quantization_method == QuantizationMethod.SYMMETRIC:
-            activation_quant_cfg.activation_quantization_params_fn = \
+            activation_quant_cfg.set_activation_quantization_params_fn(
                 quantization_params_generation.symmetric_no_clipping_selection_min_max
+            )
         elif activation_quant_cfg.activation_quantization_method == QuantizationMethod.UNIFORM:
-            activation_quant_cfg.activation_quantization_params_fn = \
+            activation_quant_cfg.set_activation_quantization_params_fn(
                 quantization_params_generation.uniform_no_clipping_selection_min_max
+            )
+
+
+def get_activations_qparams(activation_quant_cfg: NodeActivationQuantizationConfig,
+                            nodes_prior_info: NodePriorInfo,
+                            out_stats_container: BaseStatsCollector) -> Dict[str, Union[np.ndarray, float, bool]]:
+    """
+    Compute the activations params for a given node in a graph according to a params function.
+
+    Args:
+        activation_quant_cfg: node's activation quantization configuration.
+        nodes_prior_info: Prior info collected for the node that is being quantized.
+        out_stats_container: Tensor containing output statistics of the node.
+
+    Returns:
+        The computed activation quantization params.
+    """
+    # Update quantization parameters function based on output bounds and quantization method.
+    update_activation_quantization_params_fn(activation_quant_cfg, nodes_prior_info)
+
+    # Extract and filter histogram data from the statistics container.
+    bins_values, bins_counts = get_histogram_data(activation_quant_cfg, out_stats_container)
+
+    # Retrieve the minimum and maximum values from the statistics container.
+    min_value, max_value = out_stats_container.get_min_max_values()
+
+    # Determine if the activations should be considered signed.
+    signed = determine_signedness(
+        activation_quant_cfg,
+        nodes_prior_info,
+        min_value,
+        bins_values,
+        bins_counts
+    )
 
-    return activation_quant_cfg.activation_quantization_params_fn(bins_values,
-                                                                  bins_counts,
-                                                                  activation_quant_cfg.l_p_value,
-                                                                  activation_quant_cfg.activation_n_bits,
-                                                                  min_value,
-                                                                  max_value,
-                                                                  min_threshold=activation_quant_cfg.min_threshold,
-                                                                  quant_error_method=activation_quant_cfg.activation_error_method,
-                                                                  is_signed=signed)
+    # Compute and return the activation quantization parameters.
+    return activation_quant_cfg.activation_quantization_params_fn(
+        bins_values,
+        bins_counts,
+        activation_quant_cfg.l_p_value,
+        activation_quant_cfg.activation_n_bits,
+        min_value,
+        max_value,
+        min_threshold=activation_quant_cfg.min_threshold,
+        quant_error_method=activation_quant_cfg.activation_error_method,
+        is_signed=signed
+    )

model_compression_toolkit/core/keras/keras_implementation.py

@@ -188,18 +188,30 @@ class KerasImplementation(FrameworkImplementation):
 
     def run_model_inference(self,
                             model: Any,
-                            input_list: List[Any]) -> Tuple[tf.Tensor]:
+                            input_list: List[Any],
+                            requires_grad: bool = False) -> Tuple[tf.Tensor]:
         """
-        Run the model logic on the given the inputs.
+        Runs inference on the given Keras model with the provided inputs.
+
+        This method executes the model on the given input data. If `requires_grad` is set to `False`,
+        gradients will not be computed during inference by wrapping execution in a `tf.stop_gradient()` context.
 
         Args:
-            model: Keras model.
-            input_list: List of inputs for the model.
+            model: The Keras model to execute.
+            input_list: A list of inputs for the model.
+            requires_grad: If False, prevents gradient computation (default: False).
 
         Returns:
-            The Keras model's output.
-        """
-        return model(input_list)
+            A tuple containing the model's output tensors.
+        """
+        # Prevent gradient computation if requires_grad is False
+        if requires_grad:
+            # Record operations for automatic differentiation
+            with tf.GradientTape(persistent=True, watch_accessed_variables=False) as g:
+                g.watch(input_list)
+                return model(input_list)
+        else:
+            return model(input_list)
 
     def shift_negative_correction(self,
                                   graph: Graph,
@@ -553,28 +565,24 @@ class KerasImplementation(FrameworkImplementation):
 
         Returns: The MAC count og the operation
         """
-
-        output_shape = node.output_shape
-        kernel_shape = node.get_weights_by_keys(fw_info.get_kernel_op_attributes(node.type)[0]).shape
-        output_channel_axis, input_channel_axis = fw_info.kernel_channels_mapping.get(node.type)
-
-        if node.is_match_type(Conv2D) or node.is_match_type(Conv2DTranspose):
-            # (C_out * W_out * H_out) * C_in * (W_kernel * H_kernel)
-            return np.prod([x for x in output_shape if x is not None]) * \
-                   kernel_shape[input_channel_axis] * \
-                   (kernel_shape[0] * kernel_shape[1])
-        elif node.is_match_type(DepthwiseConv2D):
-            # Depth * (W_out * H_out) * C_in * (W_kernel * H_kernel)
-            return node.framework_attr.get(DEPTH_MULTIPLIER) * \
-                   np.prod([x for x in output_shape if x is not None]) / output_shape[output_channel_axis] * \
-                   kernel_shape[input_channel_axis] * \
-                   (kernel_shape[0] * kernel_shape[1])
-        elif node.is_match_type(Dense):
-            # IN * OUT
-            return kernel_shape[0] * kernel_shape[1]
-        else:
+        kernels = fw_info.get_kernel_op_attributes(node.type)
+        if not kernels or kernels[0] is None:
             return 0
 
+        assert len(kernels) == 1
+        kernel_shape = node.get_weights_by_keys(kernels[0]).shape
+
+        if node.is_match_type(Conv2D) or node.is_match_type(Conv2DTranspose) or node.is_match_type(DepthwiseConv2D):
+            h, w = node.get_output_shapes_list()[0][-3:-1]
+            return np.prod(kernel_shape) * h * w
+
+        if node.is_match_type(Dense):
+            # IN * OUT * (all previous dims[:-1])
+            _, input_channel_axis = fw_info.kernel_channels_mapping.get(node.type)
+            return node.get_total_output_params() * kernel_shape[input_channel_axis]
+
+        return 0
+
     def apply_second_moment_correction(self,
                                        quantized_model: Any,
                                        core_config: CoreConfig,

model_compression_toolkit/core/pytorch/hessian/activation_hessian_scores_calculator_pytorch.py

@@ -72,19 +72,21 @@ class ActivationHessianScoresCalculatorPytorch(HessianScoresCalculatorPytorch):
             input_tensor.requires_grad_()
             input_tensor.retain_grad()
 
-        outputs = model(*self.input_images)
+        model_output_tensors = model(*self.input_images)
 
-        if len(outputs) != len(grad_model_outputs):  # pragma: no cover
+        if len(model_output_tensors) != len(grad_model_outputs):  # pragma: no cover
             Logger.critical(f"Mismatch in expected and actual model outputs for activation Hessian approximation. "
-                            f"Expected {len(grad_model_outputs)} outputs, received {len(outputs)}.")
+                            f"Expected {len(grad_model_outputs)} outputs, received {len(model_output_tensors)}.")
+        return model_output_tensors
 
+    def _prep_tensors_for_compute(self, model_output_tensors):
         # Extracting the intermediate activation tensors and the model real output.
         # Note that we do not allow computing Hessian for output nodes, so there shouldn't be an overlap.
         num_target_nodes = len(self.hessian_request.target_nodes)
         # Extract activation tensors of nodes for which we want to compute Hessian
-        target_activation_tensors = outputs[:num_target_nodes]
+        target_activation_tensors = model_output_tensors[:num_target_nodes]
         # Extract the model outputs
-        output_tensors = outputs[num_target_nodes:]
+        output_tensors = model_output_tensors[num_target_nodes:]
 
         # Concat outputs
         # First, we need to unfold all outputs that are given as list, to extract the actual output tensors
@@ -98,79 +100,39 @@ class ActivationHessianScoresCalculatorPytorch(HessianScoresCalculatorPytorch):
         Returns:
             List[np.ndarray]: Scores based on the approximated Hessian for the requested nodes.
         """
-        output, target_activation_tensors = self.forward_pass()
-
-        if self.hessian_request.granularity == HessianScoresGranularity.PER_TENSOR:
-            hessian_scores = self._compute_per_tensor(output, target_activation_tensors)
-        elif self.hessian_request.granularity == HessianScoresGranularity.PER_OUTPUT_CHANNEL:
-            hessian_scores = self._compute_per_channel(output, target_activation_tensors)
+        if self.hessian_request.compute_from_tensors:
+            model_output_tensors = self.input_images
         else:
-            raise NotImplementedError(f'{self.hessian_request.granularity} is not supported')    # pragma: no cover
-
-        # Convert results to list of numpy arrays
-        hessian_results = [torch_tensor_to_numpy(h) for h in hessian_scores]
-        return hessian_results
+            model_output_tensors = self.forward_pass()
+        output, target_activation_tensors = self._prep_tensors_for_compute(model_output_tensors)
 
-    def _compute_per_tensor(self, output, target_activation_tensors):
-        assert self.hessian_request.granularity == HessianScoresGranularity.PER_TENSOR
-        ipts_hessian_approx_scores = [torch.tensor([0.0], requires_grad=True, device=output.device)
+        ipts_hessian_approx_scores = [torch.tensor(0.0, requires_grad=True, device=output.device)
                                       for _ in range(len(target_activation_tensors))]
-        prev_mean_results = None
-        for j in tqdm(range(self.num_iterations_for_approximation), "Hessian random iterations"):  # Approximation iterations
-            # Getting a random vector
+
+        for j in tqdm(range(self.num_iterations_for_approximation),
+                      "Hessian random iterations"):  # Approximation iterations
             v = self._generate_random_vectors_batch(output.shape, output.device)
             f_v = torch.sum(v * output)
             for i, ipt_tensor in enumerate(target_activation_tensors):  # Per Interest point activation tensor
-                # 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,
                                        allow_unused=True)[0]
-
                 if hess_v is None:
                     # In case we have an output node, which is an interest point, but it is not differentiable,
                     # we consider its Hessian to be the initial value 0.
                     continue  # pragma: no cover
 
-                # Mean over all dims but the batch (CXHXW for conv)
-                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_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_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_approx_scores), dim=1)
-
-        # add extra dimension to preserve previous behaviour
-        ipts_hessian_approx_scores = [torch.unsqueeze(t, -1) for t in ipts_hessian_approx_scores]
-        return ipts_hessian_approx_scores
-
-    def _compute_per_channel(self, output, target_activation_tensors):
-        assert self.hessian_request.granularity == HessianScoresGranularity.PER_OUTPUT_CHANNEL
-        ipts_hessian_approx_scores = [torch.tensor(0.0, requires_grad=True, device=output.device)
-                                      for _ in range(len(target_activation_tensors))]
-
-        for j in tqdm(range(self.num_iterations_for_approximation), "Hessian random iterations"):  # Approximation iterations
-            v = self._generate_random_vectors_batch(output.shape, output.device)
-            f_v = torch.sum(v * output)
-            for i, ipt_tensor in enumerate(target_activation_tensors):  # Per Interest point activation tensor
-                hess_v = autograd.grad(outputs=f_v,
-                                       inputs=ipt_tensor,
-                                       retain_graph=True)[0]
                 hessian_approx_scores = hess_v ** 2
-                rank = len(hess_v.shape)
-                if rank > 2:
-                    hessian_approx_scores = torch.mean(hessian_approx_scores, dim=tuple(range(2, rank)))
+                num_dims = len(hess_v.shape)
+                if self.hessian_request.granularity == HessianScoresGranularity.PER_TENSOR:
+                    hessian_approx_scores = torch.sum(hessian_approx_scores, dim=tuple(range(1,num_dims))).unsqueeze(-1)
+                elif self.hessian_request.granularity == HessianScoresGranularity.PER_OUTPUT_CHANNEL and num_dims > 2:
+                    hessian_approx_scores = torch.mean(hessian_approx_scores, dim=tuple(range(2, num_dims)))
 
                 # Update node Hessian approximation mean over random iterations
                 ipts_hessian_approx_scores[i] = (j * ipts_hessian_approx_scores[i] + hessian_approx_scores) / (j + 1)
 
-        return ipts_hessian_approx_scores
+        # Convert results to list of numpy arrays
+        hessian_results = [torch_tensor_to_numpy(h) for h in ipts_hessian_approx_scores]
+        return hessian_results