feat: add MSE evaluation metric

docs/user_manual/evaluate.rst

@@ -10,7 +10,7 @@ Evaluation helps you understand how compression affects your models across diffe
 This knowledge is essential for making informed decisions about which compression techniques work best for your specific needs using two types of metrics:
 
 - **Efficiency Metrics:** Measure speed (total time, latency, throughput), memory (disk, inference, training), and energy usage (consumption, CO2 emissions).
-- **Quality Metrics:** Assess fidelity (FID, CMMD), alignment (Clip Score), diversity (PSNR, SSIM), accuracy (accuracy, precision, perplexity), and more. Custom metrics are supported.
+- **Quality Metrics:** Assess fidelity (FID, CMMD, MSE), alignment (Clip Score), diversity (PSNR, SSIM), accuracy (accuracy, precision, perplexity), and more. Custom metrics are supported.
 
 .. image:: /_static/assets/images/evaluation_agent.png
     :alt: Evaluation Agent
@@ -275,11 +275,11 @@ This is what's happening under the hood when you pass ``call_type="single"`` or
 
    * - ``y_gt``
      - Model's output first, then ground truth
-     - ``fid``, ``cmmd``, ``accuracy``, ``recall``, ``precision``
+     - ``fid``, ``cmmd``, ``mse``, ``accuracy``, ``recall``, ``precision``
 
    * - ``gt_y``
      - Ground truth first, then model's output
-     - ``fid``, ``cmmd``, ``accuracy``, ``recall``, ``precision``
+     - ``fid``, ``cmmd``, ``mse``, ``accuracy``, ``recall``, ``precision``
 
    * - ``x_gt``
      - Input data first, then ground truth
@@ -291,15 +291,15 @@ This is what's happening under the hood when you pass ``call_type="single"`` or
 
    * - ``pairwise``
      - Pairwise mode to default to ``pairwise_y_gt`` or ``pairwise_gt_y``
-     - ``psnr``, ``ssim``, ``lpips``, ``cmmd``
+     - ``psnr``, ``ssim``, ``lpips``, ``cmmd``, ``mse``
 
    * - ``pairwise_y_gt``
      - Base model's output first, then subsequent model's output
-     -  ``psnr``, ``ssim``, ``lpips``, ``cmmd``
+     -  ``psnr``, ``ssim``, ``lpips``, ``cmmd``, ``mse``
 
    * - ``pairwise_gt_y``
      - Subsequent model's output first, then base model's output
-     - ``psnr``, ``ssim``, ``lpips``, ``cmmd``
+     - ``psnr``, ``ssim``, ``lpips``, ``cmmd``, ``mse``
 
    * - ``y``
      - Only the output is used, the metric has an internal dataset

src/pruna/evaluation/metrics/__init__.py

@@ -22,6 +22,7 @@
 from pruna.evaluation.metrics.metric_evalharness import LMEvalMetric
 from pruna.evaluation.metrics.metric_memory import DiskMemoryMetric, InferenceMemoryMetric, TrainingMemoryMetric
 from pruna.evaluation.metrics.metric_model_architecture import TotalMACsMetric, TotalParamsMetric
+from pruna.evaluation.metrics.metric_mse import MSE
 from pruna.evaluation.metrics.metric_pairwise_clip import PairwiseClipScore
 from pruna.evaluation.metrics.metric_rapiddata import RapidataMetric as RapidataMetric
 from pruna.evaluation.metrics.metric_sharpness import SharpnessMetric
@@ -42,6 +43,7 @@
     "TotalMACsMetric",
     "PairwiseClipScore",
     "CMMD",
+    "MSE",
     "DinoScore",
     "SharpnessMetric",
     "AestheticLAION",

src/pruna/evaluation/metrics/metric_mse.py

@@ -0,0 +1,137 @@
+# Copyright 2025 - Pruna AI GmbH. 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 __future__ import annotations
+
+from typing import Any, List
+
+import torch
+
+from pruna.engine.utils import device_to_string
+from pruna.evaluation.metrics.metric_stateful import StatefulMetric
+from pruna.evaluation.metrics.registry import MetricRegistry
+from pruna.evaluation.metrics.result import MetricResult
+from pruna.evaluation.metrics.utils import SINGLE, get_call_type_for_single_metric, metric_data_processor
+from pruna.logging.logger import pruna_logger
+
+METRIC_MSE = "mse"
+
+
+@MetricRegistry.register(METRIC_MSE)
+class MSE(StatefulMetric):
+    """
+    Calculates the Mean Squared Error (MSE) between two tensors.
+
+    The MSE is the average of the element-wise squared differences between two tensors.
+    When comparing images, this is equivalent to overlaying the two images and averaging the squared
+    pixel differences: identical images give an MSE of ``0.0`` and larger values indicate that the
+    tensors are further apart.
+
+    The error is accumulated over all batches seen during evaluation, so the reported value is the
+    mean squared error over every element of every batch.
+
+    Parameters
+    ----------
+    *args : Any
+        Additional arguments to pass to the StatefulMetric constructor.
+    device : str | torch.device | None, optional
+        The device to be used, e.g., 'cuda' or 'cpu'. Default is None.
+        If None, the best available device will be used.
+    call_type : str
+        The call type to use for the metric. Determines which two tensors are compared,
+        e.g. the model outputs against the ground truth. Defaults to single mode.
+    **kwargs : Any
+        Additional keyword arguments to pass to the StatefulMetric constructor.
+    """
+
+    sum_squared_error: torch.Tensor
+    n_observations: torch.Tensor
+    metric_name: str = METRIC_MSE
+    higher_is_better: bool = False
+    default_call_type: str = "gt_y"
+
+    def __init__(
+        self,
+        *args,
+        device: str | torch.device | None = None,
+        call_type: str = SINGLE,
+        **kwargs,
+    ) -> None:
+        super().__init__(device=device)
+        self.call_type = get_call_type_for_single_metric(call_type, self.default_call_type)
+
+        self.add_state("sum_squared_error", torch.tensor(0.0, device=self.device))
+        self.add_state("n_observations", torch.tensor(0, device=self.device))
+
+    @torch.no_grad()
+    def update(self, x: List[Any] | torch.Tensor, gt: List[Any] | torch.Tensor, outputs: Any) -> None:
+        """
+        Update the metric with new batch data.
+
+        The two tensors selected by the call type are compared element-wise and the sum of their
+        squared differences is accumulated together with the number of compared elements.
+
+        Parameters
+        ----------
+        x : List[Any] | torch.Tensor
+            The input data.
+        gt : List[Any] | torch.Tensor
+            The ground truth / cached images.
+        outputs : Any
+            The output images.
+        """
+        inputs = metric_data_processor(x, gt, outputs, self.call_type, self.device)
+        a = torch.as_tensor(inputs[0], dtype=torch.float32, device=self.device)
+        b = torch.as_tensor(inputs[1], dtype=torch.float32, device=self.device)
+
+        if a.shape != b.shape:
+            msg = f"MSE requires matching tensor shapes, got {tuple(a.shape)} and {tuple(b.shape)}."
+            pruna_logger.error(msg)
+            raise ValueError(msg)
+
+        self.sum_squared_error += torch.sum((a - b) ** 2)
+        self.n_observations += a.numel()
+
+    def compute(self) -> MetricResult:
+        """
+        Compute the mean squared error over all accumulated batches.
+
+        Returns
+        -------
+        MetricResult
+            The computed MSE metric result.
+        """
+        if self.n_observations == 0:
+            pruna_logger.warning("MSE receives no samples, returning NaN.")
+            return MetricResult(self.metric_name, self.__dict__.copy(), float("nan"))
+
+        result = (self.sum_squared_error / self.n_observations).item()
+        return MetricResult(self.metric_name, self.__dict__.copy(), result)
+
+    def move_to_device(self, device: str | torch.device) -> None:
+        """
+        Move the metric to a specific device.
+
+        Parameters
+        ----------
+        device : str | torch.device
+            The device to move the metric to.
+        """
+        if not self.is_device_supported(device):
+            raise ValueError(
+                f"Metric {self.metric_name} does not support device {device}. Must be one of {self.runs_on}."
+            )
+        self.device = device_to_string(device)
+        self.sum_squared_error = self.sum_squared_error.to(device)
+        self.n_observations = self.n_observations.to(device)

tests/evaluation/test_mse.py

@@ -0,0 +1,130 @@
+import math
+import pytest
+import torch
+
+from pruna.data.pruna_datamodule import PrunaDataModule
+from pruna.evaluation.metrics import MetricRegistry
+from pruna.evaluation.metrics.metric_mse import MSE
+
+
+@pytest.mark.parametrize(
+    "datamodule_fixture, device",
+    [
+        pytest.param("LAION256", "cpu", marks=pytest.mark.cpu),
+        pytest.param("LAION256", "cuda", marks=pytest.mark.cuda),
+    ],
+    indirect=["datamodule_fixture"],
+)
+def test_mse_on_pruna_data(datamodule_fixture: PrunaDataModule, device: str) -> None:
+    """Test the MSE on real images: zero for identical images and positive for different ones."""
+    metric = MSE(device=device)
+    dataloader = datamodule_fixture.val_dataloader()
+    dataloader_iter = iter(dataloader)
+
+    x, gt1 = next(dataloader_iter)
+    _, gt2 = next(dataloader_iter)
+
+    # Identical images overlap perfectly, so their MSE is 0.0
+    metric.update(x, gt1, gt1)
+    assert metric.compute().result == pytest.approx(0.0)
+
+    # Two different images have a positive pixel-wise error
+    metric.reset()
+    metric.update(x, gt1, gt2)
+    assert metric.compute().result > 0.0
+
+
+@pytest.mark.parametrize(
+    "device",
+    [
+        pytest.param("cpu", marks=pytest.mark.cpu),
+        pytest.param("cuda", marks=pytest.mark.cuda),
+    ],
+)
+def test_mse_identical_tensors(device: str) -> None:
+    """Test that the MSE of two identical tensors is zero."""
+    image = torch.rand(1, 3, 16, 16)
+
+    metric = MSE(device=device)
+    metric.update(image, image, image)
+
+    assert metric.compute().result == pytest.approx(0.0)
+
+
+@pytest.mark.parametrize(
+    "device",
+    [
+        pytest.param("cpu", marks=pytest.mark.cpu),
+        pytest.param("cuda", marks=pytest.mark.cuda),
+    ],
+)
+def test_mse_known_value(device: str) -> None:
+    """Test the MSE against a simple pre-computed example."""
+    gt = torch.zeros(2, 2)
+    outputs = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
+
+    metric = MSE(device=device)
+    metric.update(gt, gt, outputs)
+
+    assert metric.compute().result == pytest.approx(7.5)
+
+
+@pytest.mark.parametrize(
+    "device",
+    [
+        pytest.param("cpu", marks=pytest.mark.cpu),
+        pytest.param("cuda", marks=pytest.mark.cuda),
+    ],
+)
+def test_mse_accumulates_and_resets(device: str) -> None:
+    """Test that the MSE is pooled across batches and that reset clears the state."""
+    gt = torch.zeros(1, 3, 4, 4) # Ground-truth of 0.0, 1*3*4*4 = 48 elements per batch
+    outputs_batch_one = torch.ones_like(gt)  # Squared error of 1.0 per element
+    outputs_batch_two = torch.full_like(gt, 3.0)  # Squared error of 9.0 per element
+
+    metric = MSE(device=device)
+    metric.update(gt, gt, outputs_batch_one)
+    metric.update(gt, gt, outputs_batch_two)
+
+    # The pooled mean is (48 * 1.0 + 48 * 9.0) / (48 + 48) = 5.0
+    assert metric.compute().result == pytest.approx(5.0)
+
+    metric.reset()
+    metric.update(gt, gt, outputs_batch_one)
+    assert metric.compute().result == pytest.approx(1.0)
+
+
+@pytest.mark.cpu
+def test_mse_no_update() -> None:
+    """Test that computing without any update explicitly returns NaN."""
+    metric = MSE(device="cpu")
+
+    assert math.isnan(metric.compute().result)
+
+
+@pytest.mark.cpu
+def test_mse_shape_mismatch_raises() -> None:
+    """Test that comparing tensors of different shapes raises a ValueError."""
+    gt = torch.zeros(1, 3, 8, 8)
+    outputs = torch.zeros(1, 3, 4, 4)
+
+    metric = MSE(device="cpu")
+    with pytest.raises(ValueError):
+        metric.update(gt, gt, outputs)
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize(
+    "call_type, expected_call_type",
+    [
+        ("single", "gt_y"),
+        ("pairwise", "pairwise_gt_y"),
+    ],
+)
+def test_mse_call_type(call_type: str, expected_call_type: str) -> None:
+    """Test that the MSE exposes the expected attributes and resolves its call type."""
+    metric = MSE(call_type=call_type, device="cpu")
+
+    assert metric.metric_name == "mse"
+    assert metric.higher_is_better is False
+    assert metric.call_type == expected_call_type