Add ReportResultsComparison

Author: xadupre

_doc/api/reference/index.rst

@@ -15,6 +15,7 @@ onnx_diagnostic.reference
     evaluator
     quantized_tensor
     ort_evaluator
+    report_results_comparison
     torch_evaluator
 
 ExtendedReferenceEvaluator
@@ -29,6 +30,12 @@ OnnxruntimeEvaluator
 .. autoclass:: onnx_diagnostic.reference.OnnxruntimeEvaluator
     :members:
 
+ReportResultsComparison
+++++++++++++++++++
+
+.. autoclass:: onnx_diagnostic.reference.ReportResultsComparison
+    :members:
+
 TorchOnnxEvaluator
 ++++++++++++++++++
 

_doc/api/reference/report_results_comparison.rst

@@ -0,0 +1,8 @@
+
+onnx_diagnostic.reference.report_results_comparison
+===================================================
+
+.. automodule:: onnx_diagnostic.reference.report_results_comparison
+    :members:
+    :no-undoc-members:
+    :exclude-members: ReportResultsComparison

_unittests/ut_reference/test_torch_onnx_evaluator.py

@@ -4,10 +4,14 @@
 import onnx.helper as oh
 import onnx.numpy_helper as onh
 import torch
-from onnx_diagnostic.ext_test_case import ExtTestCase, ignore_warnings
+from onnx_diagnostic.ext_test_case import ExtTestCase, ignore_warnings, hide_stdout
 from onnx_diagnostic.helpers.onnx_helper import from_array_extended
 from onnx_diagnostic.helpers.torch_helper import onnx_dtype_to_torch_dtype
-from onnx_diagnostic.reference import ExtendedReferenceEvaluator, TorchOnnxEvaluator
+from onnx_diagnostic.reference import (
+    ExtendedReferenceEvaluator,
+    TorchOnnxEvaluator,
+    ReportResultsComparison,
+)
 from onnx_diagnostic.reference.torch_ops import OpRunKernel, OpRunTensor
 from onnx_diagnostic.reference.torch_evaluator import get_kernels
 
@@ -1471,6 +1475,35 @@ def run(self, x, scale, bias=None):
         self.assertEqualAny(expected, got, atol=1e-3)
         self.assertEqual([1], LayerNormalizationOrt._shared)
 
+    @hide_stdout()
+    def test_report_results_comparison(self):
+        model = oh.make_model(
+            oh.make_graph(
+                [
+                    oh.make_node("Cos", ["X"], ["nx"]),
+                    oh.make_node("Sin", ["nx"], ["t"]),
+                    oh.make_node("Exp", ["t"], ["u"]),
+                    oh.make_node("Log", ["u"], ["uZ"]),
+                    oh.make_node("Erf", ["uZ"], ["Z"]),
+                ],
+                "dummy",
+                [oh.make_tensor_value_info("X", TFLOAT, ["a", "b"])],
+                [oh.make_tensor_value_info("Z", TFLOAT, ["a", "b"])],
+            ),
+            ir_version=9,
+            opset_imports=[oh.make_opsetid("", 18)],
+        )
+        x = torch.rand(5, 6, dtype=torch.float32)
+        onnx.checker.check_model(model)
+        cmp = ReportResultsComparison(dict(r_x=x, r_cos=x.cos(), r_exp=x.cos().sin().exp()))
+        cmp.clear()
+        feeds = dict(zip([i.name for i in model.graph.input], (x,)))
+        rt = TorchOnnxEvaluator(model, verbose=10)
+        rt.run(None, feeds, report_cmp=cmp)
+        d = {k: d["abs"] for k, d in cmp.value.items()}
+        self.assertEqual(d["nx", "r_cos"], 0)
+        self.assertEqual(d["u", "r_exp"], 0)
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

onnx_diagnostic/reference/__init__.py

@@ -1,3 +1,4 @@
 from .evaluator import ExtendedReferenceEvaluator
 from .ort_evaluator import OnnxruntimeEvaluator
 from .torch_evaluator import TorchOnnxEvaluator
+from .report_results_comparison import ReportResultsComparison

onnx_diagnostic/reference/report_results_comparison.py

@@ -0,0 +1,75 @@
+from typing import Dict, List, Tuple, Union
+
+
+ReportKeyType = Union[str, Tuple[Union[int, str], ...]]
+
+
+class ReportResultsComparison:
+    """
+    Holds tensors a runtime can use as a reference to compare
+    intermediate results.
+    See :meth:`onnx_diagnostic.reference.TorchOnnxEvalutor.run`.
+
+    :param tensors: tensor
+    """
+
+    def __init__(self, tensors: Dict[ReportKeyType, "torch.Tensor"]):  # noqa: F821
+        from ..helpers.onnx_helper import dtype_to_tensor_dtype
+        from ..helpers import max_diff
+
+        self.dtype_to_tensor_dtype = dtype_to_tensor_dtype
+        self.max_diff = max_diff
+        self.tensors = tensors
+        self._build_mapping()
+
+    def key(self, tensor: "torch.Tensor") -> ReportKeyType:  # noqa: F821
+        "Returns a key for a tensor, (onnx dtype, shape)."
+        return self.dtype_to_tensor_dtype(tensor.dtype), tuple(map(int, tensor.shape))
+
+    def _build_mapping(self):
+        mapping = {}
+        for k, v in self.tensors.items():
+            key = self.key(v)
+            if key not in mapping:
+                mapping[key] = []
+            mapping[key].append(k)
+        self.mapping = mapping
+        self.clear()
+
+    def clear(self):
+        """Clears the last report."""
+        self.report_cmp = {}
+
+    @property
+    def value(self) -> Dict[Tuple[str, ReportKeyType], Dict[str, Union[float, str]]]:
+        "Returns the report."
+        return self.report_cmp
+
+    def report(
+        self, outputs: Dict[str, "torch.Tensor"]  # noqa: F821
+    ) -> List[Tuple[str, ReportKeyType]]:
+        """
+        For every tensor in outputs, compares it to every tensor held by
+        this class if it shares the same type and shape. The function returns
+        the results of the comparison. The function also collects the results
+        into a dictionary the user can retrieve later.
+        """
+        res = []
+        for name, tensor in outputs.items():
+            key = self.key(tensor)
+            if key not in self.mapping:
+                continue
+            cache = {}
+            for held_key in self.mapping[key]:
+                t2 = self.tensors[held_key]
+                if hasattr(t2, "device") and hasattr(tensor, "device"):
+                    if t2.device in cache:
+                        t = cache[t2.device]
+                    else:
+                        cache[t2.device] = t = tensor.to(t2.device)
+                    diff = self.max_diff(t, t2)
+                else:
+                    diff = self.max_diff(tensor, t2)
+                res.append((name, held_key, diff))
+                self.report_cmp[name, held_key] = diff
+        return res

onnx_diagnostic/reference/torch_evaluator.py

@@ -5,6 +5,7 @@
 import torch
 from ..helpers.torch_helper import to_tensor
 from ..torch_onnx.runtime_info import first_used_last_used, RuntimeValue
+from .report_results_comparison import ReportResultsComparison
 from . import torch_ops
 
 
@@ -455,12 +456,17 @@ def run(
         self,
         outputs: Optional[List[str]],
         feeds: Union[Dict[str, torch.Tensor], Dict[str, np.ndarray]],
+        report_cmp: Optional[ReportResultsComparison] = None,
     ) -> Union[List[Optional[torch.Tensor]], List[Optional[np.ndarray]]]:
         """
         Runs the ONNX model.
 
         :param outputs: outputs required
         :param feeds: inputs
+        :param report_cmp: used as a reference,
+            every intermediate results is compare to every existing one,
+            if not empty, it is an instance of
+            :class:`onnx_diagnostic.reference.ReportResultsComparison`
         :return: output tensors.
         """
         use_numpy = any(isinstance(t, np.ndarray) for t in feeds.values())
@@ -532,6 +538,21 @@ def run(
                             f"+R {kernel.output[0]}: "
                             f"{self.runtime_info[kernel.output[0]].string_type()}"
                         )
+                if report_cmp:
+                    reported = report_cmp.report(
+                        dict(
+                            zip(
+                                kernel.output,
+                                (
+                                    tuple(r.tensor for r in res)
+                                    if isinstance(res, tuple)
+                                    else (res.tensor,)
+                                ),
+                            )
+                        )
+                    )
+                    if self.verbose > 1:
+                        print(f"  -- report {len(reported)} comparisons")
 
             # free intermediate results
             for name in self.last_used[it]: