Add SNR numerical comparator

Summary: This PR introduces the SNR (Signal-to-Noise Ratio) comparator class, which extends the numerical comparison framework to evaluate model accuracy by comparing the SNR between two inputs. The comparator calculates the signal power and noise power from the two inputs, then computes the SNR using the formula SNR = 10 * log10(original_power / error_power).

Differential Revision: D77159515

Author: Juntian Liu

devtools/inspector/_inspector_utils.py

@@ -720,4 +720,8 @@ def convert_to_float_tensor(input_data: Any) -> torch.Tensor:
             f"Cannot convert value of type {type(input_data)} to a tensor: {e}"
         )
     input_tensor = input_tensor.detach().cpu().double()
+
+    # Convert NaN to 0.0
+    if torch.isnan(input_tensor).any():
+        input_tensor = torch.nan_to_num(input_tensor)
     return input_tensor

devtools/inspector/numerical_comparator/TARGETS

@@ -27,11 +27,21 @@ python_library(
     ],
 )
 
+python_library(
+    name = "snr_numerical_comparator",
+    srcs = ["snr_numerical_comparator.py"],
+    deps = [
+        "//executorch/devtools/inspector/numerical_comparator:numerical_comparator_base",
+        "//executorch/devtools/inspector:inspector_utils",
+    ],
+)
+
 python_library(
     name = "lib",
     srcs = ["__init__.py"],
     deps = [
         ":l1_numerical_comparator",
         ":mse_numerical_comparator",
+        ":snr_numerical_comparator",
     ],
 )

devtools/inspector/numerical_comparator/__init__.py

@@ -13,5 +13,9 @@
     MSEComparator,
 )
 
+from executorch.devtools.inspector.numerical_comparator.snr_numerical_comparator import (
+    SNRComparator,
+)
+
 
-__all__ = ["L1Comparator", "MSEComparator"]
+__all__ = ["L1Comparator", "MSEComparator", "SNRComparator"]

devtools/inspector/numerical_comparator/l1_numerical_comparator.py

@@ -19,9 +19,6 @@ def compare(self, a: Any, b: Any) -> float:
 
         t_a = convert_to_float_tensor(a)
         t_b = convert_to_float_tensor(b)
-        if torch.isnan(t_a).any() or torch.isnan(t_b).any():
-            t_a = torch.nan_to_num(t_a)
-            t_b = torch.nan_to_num(t_b)
 
         try:
             res = torch.abs(t_a - t_b).sum().item()

devtools/inspector/numerical_comparator/mse_numerical_comparator.py

@@ -19,9 +19,6 @@ def compare(self, a: Any, b: Any) -> float:
 
         t_a = convert_to_float_tensor(a)
         t_b = convert_to_float_tensor(b)
-        if torch.isnan(t_a).any() or torch.isnan(t_b).any():
-            t_a = torch.nan_to_num(t_a)
-            t_b = torch.nan_to_num(t_b)
 
         try:
             res = float(torch.mean(torch.square(t_a - t_b)))

devtools/inspector/numerical_comparator/snr_numerical_comparator.py

@@ -0,0 +1,39 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+from typing import Any
+
+import torch
+from executorch.devtools.inspector._inspector_utils import convert_to_float_tensor
+from executorch.devtools.inspector.numerical_comparator.numerical_comparator_base import (
+    NumericalComparatorBase,
+)
+
+
+class SNRComparator(NumericalComparatorBase):
+    def compare(self, a: Any, b: Any) -> float:
+        """
+        Compare the Signal-to-Noise Ratio (SNR) between two inputs
+        Formula: SNR = 10 * log10(original_power / error_power)
+        """
+
+        t_a = convert_to_float_tensor(a)
+        t_b = convert_to_float_tensor(b)
+
+        # Calculate the signal power and noise power
+        original_power = torch.mean(torch.pow(t_a, 2))
+        try:
+            error = t_a - t_b
+            error_power = torch.mean(torch.pow(error, 2))
+        except Exception as e:
+            raise ValueError(
+                f"Error computing SNR difference between tensors: {str(e)}"
+            )
+
+        # Calculate SNR
+        snr = 10 * torch.log10(original_power / error_power)
+        return snr.item()

devtools/inspector/tests/TARGETS

@@ -70,6 +70,14 @@ python_unittest(
     ],
 )
 
+python_unittest(
+    name = "snr_comparator_test",
+    srcs = ["snr_comparator_test.py"],
+    deps = [
+        "//executorch/devtools/inspector/numerical_comparator:lib",
+    ],
+)
+
 python_library(
     name = "inspector_test_utils",
     srcs = [

devtools/inspector/tests/snr_comparator_test.py

@@ -0,0 +1,62 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import math
+import unittest
+
+import torch
+
+from executorch.devtools.inspector.numerical_comparator import SNRComparator
+
+
+class TestSNRComparator(unittest.TestCase):
+    snr_comparator = SNRComparator()
+
+    def test_identical_tensors(self):
+        # identical tensors --> error_power == 0 --> SNR is inf
+        a = torch.tensor([[10, 4], [3, 4]])
+        b = torch.tensor([[10, 4], [3, 4]])
+        result = self.snr_comparator.compare(a, b)
+        self.assertTrue(math.isinf(result) and result > 0)
+
+    def test_scalar(self):
+        # original_power == 1, error_power == 1 --> SNR = 10 * log10(1/1) = 0
+        a = 1
+        b = 2
+        result = self.snr_comparator.compare(a, b)
+        self.assertAlmostEqual(result, 0.0)
+
+    def test_with_nans_replaced_with_zero(self):
+        a = torch.tensor([float("nan"), 1.0])
+        b = torch.tensor([0.0, 1.0])
+        result = self.snr_comparator.compare(a, b)
+        self.assertTrue(math.isinf(result) and result > 0)
+
+    def test_shape_mismatch_raises_exception(self):
+        a = torch.tensor([1, 2, -1])
+        b = torch.tensor([1, 1, -3, 4])
+        with self.assertRaises(ValueError):
+            self.snr_comparator.compare(a, b)
+
+    def test_2D_tensors(self):
+        # original_power = mean([16, 81, 36, 16]) = 37.25
+        # error = a - b = [3, 7, 3, -1] squared = [9, 49, 9, 1] mean = 68/4 = 17.0
+        # SNR = 10 * log10(37.25/17.0)
+        a = torch.tensor([[4, 9], [6, 4]])
+        b = torch.tensor([[1, 2], [3, 5]])
+        expected = 10 * math.log10(37.25 / 17.0)
+        result = self.snr_comparator.compare(a, b)
+        self.assertAlmostEqual(result, expected)
+
+    def test_list_of_tensors(self):
+        # original_power = mean(4, 16, 25, 4]) = 12.25
+        # error = a - b = [1, 2, 2, -3] squared = [1, 4, 4, 9] mean = 18/4 = 4.5
+        # SNR = 10 * log10(37.25/17.0)
+        a = [torch.tensor([2, 4]), torch.tensor([5, 2])]
+        b = [torch.tensor([1, 2]), torch.tensor([3, 5])]
+        expected = 10 * math.log10(12.25 / 4.5)
+        result = self.snr_comparator.compare(a, b)
+        self.assertAlmostEqual(result, expected)