mct-nightly 2.2.0.20241006.532__tar.gz → 2.2.0.20241008.450__tar.gz

{mct-nightly-2.2.0.20241006.532 → mct-nightly-2.2.0.20241008.450}/PKG-INFO

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

{mct-nightly-2.2.0.20241006.532 → mct-nightly-2.2.0.20241008.450}/mct_nightly.egg-info/PKG-INFO

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

{mct-nightly-2.2.0.20241006.532 → mct-nightly-2.2.0.20241008.450}/model_compression_toolkit/__init__.py

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

{mct-nightly-2.2.0.20241006.532 → mct-nightly-2.2.0.20241008.450}/model_compression_toolkit/core/common/hessian/__init__.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from model_compression_toolkit.core.common.hessian.hessian_scores_request import HessianScoresRequest, HessianMode, HessianScoresGranularity
+from model_compression_toolkit.core.common.hessian.hessian_scores_request import (
+    HessianScoresRequest, HessianMode, HessianScoresGranularity, HessianEstimationDistribution
+)
 from model_compression_toolkit.core.common.hessian.hessian_info_service import HessianInfoService
 import model_compression_toolkit.core.common.hessian.hessian_info_utils as hessian_utils

{mct-nightly-2.2.0.20241006.532 → mct-nightly-2.2.0.20241008.450}/model_compression_toolkit/core/common/hessian/hessian_info_service.py

@@ -12,16 +12,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+import hashlib
 
 import numpy as np
 from functools import partial
 from tqdm import tqdm
-from typing import Callable, List, Dict, Any, Tuple
+from typing import Callable, List, Dict, Any, Tuple, TYPE_CHECKING
 
 from model_compression_toolkit.constants import HESSIAN_NUM_ITERATIONS
 from model_compression_toolkit.core.common.hessian.hessian_scores_request import HessianScoresRequest, \
     HessianScoresGranularity, HessianMode
 from model_compression_toolkit.logger import Logger
+if TYPE_CHECKING:    # pragma: no cover
+    from model_compression_toolkit.core.common import BaseNode
 
 
 class HessianInfoService:
@@ -228,6 +231,61 @@ class HessianInfoService:
         return next_iter_remain_samples if next_iter_remain_samples is not None and len(next_iter_remain_samples) > 0 \
         and len(next_iter_remain_samples[0]) > 0 else None
 
+    def compute_trackable_per_sample_hessian(self,
+                                             hessian_scores_request: HessianScoresRequest,
+                                             inputs_batch: List[np.ndarray]) -> Dict[str, Dict['BaseNode', np.ndarray]]:
+        """
+        Compute hessian score per image hash. We compute the score directly for images rather than via data generator,
+        as data generator might yield different images each time, depending on how it was defined,
+
+        Args:
+            hessian_scores_request: hessian scores request
+            inputs_batch: a list containing a batch of inputs.
+
+        Returns:
+            A dict of Hessian scores per image hash per layer {image hash: {layer: score}}
+        """
+        topo_sorted_nodes_names = [x.name for x in self.graph.get_topo_sorted_nodes()]
+        hessian_scores_request.target_nodes.sort(key=lambda x: topo_sorted_nodes_names.index(x.name))
+
+        hessian_score_by_image_hash = {}
+
+        if not inputs_batch or not isinstance(inputs_batch, list):
+            raise TypeError('Expected a non-empty list of inputs')    # pragma: no cover
+        if len(inputs_batch) > 1:
+            raise NotImplementedError('Per-sample hessian computation is not supported for networks with multiple inputs')    # pragma: no cover
+
+        # Get the framework-specific calculator Hessian-approximation scores
+        fw_hessian_calculator = self.fw_impl.get_hessian_scores_calculator(graph=self.graph,
+                                                                           input_images=inputs_batch,
+                                                                           hessian_scores_request=hessian_scores_request,
+                                                                           num_iterations_for_approximation=self.num_iterations_for_approximation)
+        hessian_scores = fw_hessian_calculator.compute()
+        for i in range(inputs_batch[0].shape[0]):
+            img_hash = self.calc_image_hash(inputs_batch[0][i])
+            hessian_score_by_image_hash[img_hash] = {
+                node: score[i] for node, score in zip(hessian_scores_request.target_nodes, hessian_scores)
+            }
+
+        return hessian_score_by_image_hash
+
+    @staticmethod
+    def calc_image_hash(image):
+        """
+        Calculates hash for an input image.
+
+        Args:
+            image: input 3d image (without batch).
+
+        Returns:
+            Image hash.
+
+        """
+        if not len(image.shape) == 3:    # pragma: no cover
+            raise ValueError(f'Expected 3d image (without batch) for image hash calculation, got {len(image.shape)}')
+        image_bytes = image.astype(np.float32).tobytes()
+        return hashlib.md5(image_bytes).hexdigest()
+
     def fetch_hessian(self,
                       hessian_scores_request: HessianScoresRequest,
                       required_size: int,
@@ -248,7 +306,7 @@ class HessianInfoService:
             OC for per-output-channel when the requested node has OC output-channels, etc.)
         """
 
-        if len(hessian_scores_request.target_nodes) == 0:
+        if len(hessian_scores_request.target_nodes) == 0:    # pragma: no cover
             return []
 
         if required_size == 0:

{mct-nightly-2.2.0.20241006.532 → mct-nightly-2.2.0.20241008.450}/model_compression_toolkit/core/common/hessian/hessian_scores_request.py

@@ -40,6 +40,14 @@ class HessianScoresGranularity(Enum):
     PER_TENSOR = 2
 
 
+class HessianEstimationDistribution(str, Enum):
+    """
+    Distribution for Hutchinson estimator random vector
+    """
+    GAUSSIAN = 'gaussian'
+    RADEMACHER = 'rademacher'
+
+
 class HessianScoresRequest:
     """
     Request configuration for the Hessian-approximation scores.
@@ -53,7 +61,8 @@ class HessianScoresRequest:
     def __init__(self,
                  mode: HessianMode,
                  granularity: HessianScoresGranularity,
-                 target_nodes: List):
+                 target_nodes: List,
+                 distribution: HessianEstimationDistribution = HessianEstimationDistribution.GAUSSIAN):
         """
         Attributes:
             mode (HessianMode): Mode of Hessian-approximation score (w.r.t weights or activations).
@@ -64,6 +73,7 @@ class HessianScoresRequest:
         self.mode = mode  # w.r.t activations or weights
         self.granularity = granularity  # per element, per layer, per channel
         self.target_nodes = target_nodes
+        self.distribution = distribution
 
     def __eq__(self, other):
         # Checks if the other object is an instance of HessianScoresRequest
@@ -71,9 +81,10 @@ class HessianScoresRequest:
         return isinstance(other, HessianScoresRequest) and \
                self.mode == other.mode and \
                self.granularity == other.granularity and \
-               self.target_nodes == other.target_nodes
+               self.target_nodes == other.target_nodes and \
+               self.distribution == other.distribution
 
     def __hash__(self):
         # Computes the hash based on the attributes.
         # The use of a tuple here ensures that the hash is influenced by all the attributes.
-        return hash((self.mode, self.granularity, tuple(self.target_nodes)))
+        return hash((self.mode, self.granularity, tuple(self.target_nodes), self.distribution))

mct-nightly-2.2.0.20241008.450/model_compression_toolkit/core/pytorch/hessian/activation_hessian_scores_calculator_pytorch.py

@@ -0,0 +1,201 @@
+# Copyright 2023 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from typing import List
+
+from torch import autograd
+from tqdm import tqdm
+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 (HessianScoresRequest, HessianScoresGranularity,
+                                                           HessianEstimationDistribution)
+from model_compression_toolkit.core.pytorch.back2framework.float_model_builder import FloatPyTorchModelBuilder
+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 ActivationHessianScoresCalculatorPytorch(HessianScoresCalculatorPytorch):
+    """
+    Pytorch implementation of the Hessian approximation scores Calculator for activations.
+    """
+    def __init__(self,
+                 graph: Graph,
+                 input_images: List[torch.Tensor],
+                 fw_impl,
+                 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 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(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 forward_pass(self):
+        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:
+            Logger.critical("Activation Hessian approximation cannot be computed for model outputs. "
+                            "Exclude output nodes from Hessian request targets.")
+
+        grad_model_outputs = self.hessian_request.target_nodes + model_output_nodes
+        model, _ = FloatPyTorchModelBuilder(graph=self.graph, append2output=grad_model_outputs).build_model()
+        model.eval()
+
+        # Run model inference
+        # Set inputs to track gradients during inference
+        for input_tensor in self.input_images:
+            input_tensor.requires_grad_()
+            input_tensor.retain_grad()
+
+        outputs = model(*self.input_images)
+
+        if len(outputs) != 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)}.")
+
+        # 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]
+        # Extract the model outputs
+        output_tensors = outputs[num_target_nodes:]
+        device = output_tensors[0].device
+
+        # Concat outputs
+        # First, we need to unfold all outputs that are given as list, to extract the actual output tensors
+        output = self.concat_tensors(output_tensors)
+        return output, target_activation_tensors
+
+    def _generate_random_vectors_batch(self, shape: tuple, distribution: HessianEstimationDistribution,
+                                       device: torch.device) -> torch.Tensor:
+        """
+        Generate a batch of random vectors for Hutchinson estimation
+
+        Args:
+            shape: target shape
+            distribution: distribution to sample from
+            device: target device
+
+        Returns:
+            Random tensor
+        """
+        if distribution == HessianEstimationDistribution.GAUSSIAN:
+            return torch.randn(shape, device=device)
+
+        if distribution == HessianEstimationDistribution.RADEMACHER:
+            v = torch.randint(high=2, size=shape, device=device)
+            v[v == 0] = -1
+            return v
+
+        raise ValueError(f'Unknown distribution {distribution}')    # pragma: no cover
+
+    def compute(self) -> List[np.ndarray]:
+        """
+        Compute the scores that are based on the approximation of the Hessian w.r.t the requested target nodes' activations.
+
+        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)
+        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
+
+    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)
+                                      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 with normal distribution
+            v = self._generate_random_vectors_batch(output.shape, self.hessian_request.distribution, 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, self.hessian_request.distribution, 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)))
+
+                # 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

{mct-nightly-2.2.0.20241006.532 → mct-nightly-2.2.0.20241008.450}/model_compression_toolkit/gptq/common/gptq_config.py

@@ -17,6 +17,7 @@ from enum import Enum
 from typing import Callable, Any, Dict, Optional
 
 from model_compression_toolkit.constants import GPTQ_HESSIAN_NUM_SAMPLES, ACT_HESSIAN_DEFAULT_BATCH_SIZE
+from model_compression_toolkit.core.common.hessian import HessianScoresGranularity, HessianEstimationDistribution
 from model_compression_toolkit.gptq.common.gptq_constants import REG_DEFAULT
 
 
@@ -39,17 +40,21 @@ class GPTQHessianScoresConfig:
     Configuration to use for computing the Hessian-based scores for GPTQ loss metric.
 
     Args:
-        hessians_num_samples (int): Number of samples to use for computing the Hessian-based scores.
+        hessians_num_samples (int|None): Number of samples to use for computing the Hessian-based scores.
+          If None, compute Hessian for all images.
         norm_scores (bool): Whether to normalize the returned scores of the weighted loss function (to get values between 0 and 1).
         log_norm (bool): Whether to use log normalization for the GPTQ Hessian-based scores.
         scale_log_norm (bool): Whether to scale the final vector of the Hessian-based scores.
         hessian_batch_size (int): The Hessian computation batch size. used only if using GPTQ with Hessian-based objective.
+        per_sample (bool): Whether to use per sample attention score.
     """
-    hessians_num_samples: int = GPTQ_HESSIAN_NUM_SAMPLES
+    hessians_num_samples: Optional[int] = GPTQ_HESSIAN_NUM_SAMPLES
     norm_scores: bool = True
     log_norm: bool = True
     scale_log_norm: bool = False
     hessian_batch_size: int = ACT_HESSIAN_DEFAULT_BATCH_SIZE
+    per_sample: bool = False
+    estimator_distribution: HessianEstimationDistribution = HessianEstimationDistribution.GAUSSIAN
 
 
 @dataclass

{mct-nightly-2.2.0.20241006.532 → mct-nightly-2.2.0.20241008.450}/model_compression_toolkit/gptq/common/gptq_training.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 import copy
+import hashlib
 from abc import ABC, abstractmethod
 import numpy as np
 from typing import Callable, List, Any, Dict
@@ -143,7 +144,11 @@ class GPTQTrainer(ABC):
             return np.asarray([1 / num_nodes for _ in range(num_nodes)])
 
         # Fetch hessian approximations for each target node
-        compare_point_to_hessian_approx_scores = self._fetch_hessian_approximations()
+        # TODO this smells like a potential bug. In hessian calculation target nodes are topo sorted and results are returned
+        # TODO also target nodes are replaced for reuse. Does this work correctly?
+        approximations = self._fetch_hessian_approximations(HessianScoresGranularity.PER_TENSOR)
+        compare_point_to_hessian_approx_scores = {node: score for node, score in zip(self.compare_points, approximations)}
+
         # Process the fetched hessian approximations to gather them per images
         hessian_approx_score_by_image = (
             self._process_hessian_approximations(compare_point_to_hessian_approx_scores))
@@ -172,29 +177,55 @@ class GPTQTrainer(ABC):
             # If log normalization is not enabled, return the mean of the approximations across images
             return np.mean(hessian_approx_score_by_image, axis=0)
 
-    def _fetch_hessian_approximations(self) -> Dict[BaseNode, List[List[float]]]:
+    def _compute_sample_layer_attention_scores(self, inputs_batch) -> Dict[str, Dict[BaseNode, np.ndarray]]:
+        """
+        Compute sample layer attention scores per image hash per layer.
+
+        Args:
+            inputs_batch: a list containing a batch of inputs.
+
+        Returns:
+            A dictionary with a structure {img_hash: {layer: score}}.
+
+        """
+        request = self._build_hessian_request(HessianScoresGranularity.PER_OUTPUT_CHANNEL)
+        hessian_batch_size = self.gptq_config.hessian_weights_config.hessian_batch_size
+
+        hessian_score_per_image_per_layer = {}
+        # If hessian batch is smaller than inputs batch, split it to hessian batches. If hessian batch is larger,
+        # it's currently ignored (TODO)
+        for i in range(0, inputs_batch[0].shape[0], hessian_batch_size):
+            inputs = [t[i: i+hessian_batch_size] for t in inputs_batch]
+            hessian_score_per_image_per_layer.update(
+                self.hessian_service.compute_trackable_per_sample_hessian(request, inputs)
+            )
+        for img_hash, v in hessian_score_per_image_per_layer.items():
+            hessian_score_per_image_per_layer[img_hash] = {k: t.max(axis=0) for k, t in v.items()}
+        return hessian_score_per_image_per_layer
+
+    def _fetch_hessian_approximations(self, granularity: HessianScoresGranularity) -> Dict[BaseNode, List[List[float]]]:
         """
         Fetches hessian approximations for each target node.
 
         Returns:
             Mapping of target nodes to their hessian approximations.
         """
-        approximations = {}
-        hessian_scores_request = HessianScoresRequest(
-            mode=HessianMode.ACTIVATION,
-            granularity=HessianScoresGranularity.PER_TENSOR,
-            target_nodes=self.compare_points
-        )
+        hessian_scores_request = self._build_hessian_request(granularity)
+
         node_approximations = self.hessian_service.fetch_hessian(
             hessian_scores_request=hessian_scores_request,
             required_size=self.gptq_config.hessian_weights_config.hessians_num_samples,
             batch_size=self.gptq_config.hessian_weights_config.hessian_batch_size
         )
+        return node_approximations
 
-        for i, target_node in enumerate(self.compare_points):
-            approximations[target_node] = node_approximations[i]
-
-        return approximations
+    def _build_hessian_request(self, granularity):
+        return HessianScoresRequest(
+            mode=HessianMode.ACTIVATION,
+            granularity=granularity,
+            target_nodes=self.compare_points,
+            distribution=self.gptq_config.hessian_weights_config.estimator_distribution
+        )
 
     def _process_hessian_approximations(self, approximations: Dict[BaseNode, List[List[float]]]) -> List:
         """

{mct-nightly-2.2.0.20241006.532 → mct-nightly-2.2.0.20241008.450}/model_compression_toolkit/gptq/pytorch/gptq_loss.py

@@ -13,8 +13,10 @@
 # limitations under the License.
 # ==============================================================================
 from typing import List
+
 import torch
 
+
 def mse_loss(y: torch.Tensor, x: torch.Tensor, normalized: bool = True) -> torch.Tensor:
     """
     Compute the MSE of two tensors.
@@ -25,7 +27,7 @@ def mse_loss(y: torch.Tensor, x: torch.Tensor, normalized: bool = True) -> torch
     Returns:
         The MSE of two tensors.
     """
-    loss = torch.nn.MSELoss()(x,y)
+    loss = torch.nn.MSELoss()(x, y)
     return loss / torch.mean(torch.square(x)) if normalized else loss
 
 
@@ -62,3 +64,36 @@ def multiple_tensors_mse_loss(y_list: List[torch.Tensor],
     else:
         return torch.mean(torch.stack(loss_values_list))
 
+
+def sample_layer_attention_loss(y_list: List[torch.Tensor],
+                                x_list: List[torch.Tensor],
+                                fxp_w_list,
+                                flp_w_list,
+                                act_bn_mean,
+                                act_bn_std,
+                                loss_weights: torch.Tensor) -> torch.Tensor:
+    """
+    Compute Sample Layer Attention loss between two lists of tensors.
+
+    Args:
+        y_list: First list of tensors.
+        x_list: Second list of tensors.
+        fxp_w_list, flp_w_list, act_bn_mean, act_bn_std: unused (needed to comply with the interface).
+        loss_weights: layer-sample weights tensor of shape (layers, batch)
+
+    Returns:
+        Sample Layer Attention loss (a scalar).
+    """
+    loss = 0
+    layers_mean_w = []
+
+    for i, (y, x, w) in enumerate(zip(y_list, x_list, loss_weights)):
+        norm = (y - x).pow(2).sum(1)
+        if len(norm.shape) > 1:
+            norm = norm.flatten(1).mean(1)
+        loss += torch.mean(w * norm)
+        layers_mean_w.append(w.mean())
+
+    loss = loss / torch.stack(layers_mean_w).max()
+    return loss
+