Add SNR numerical comparator

Differential Revision: D77159515

Pull Request resolved: #11859

Author: Juntian Liu

devtools/inspector/_inspector_utils.py

@@ -729,4 +729,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)