Supports error distribution in max_diff (#92)

* refactors max_diff

* updtae string_diff

* add model_statistics

* add mb

* ut

Author: xadupre

CHANGELOGS.rst

@@ -4,6 +4,8 @@ Change Logs
 0.5.0
 +++++
 
+* :pr:`92`: support errors distribution in max_diff
+* :pr:`91`: enable strings in ``guess_dynamic_shapes``
 * :pr:`88`, :pr:`89`: extends ``steal_forward`` to dump input, outputs in onnx models
 * :pr:`83`, :pr:`85`: improves the automated rewriting of control flow (test)
 

_unittests/ut_helpers/test_helper.py

@@ -127,7 +127,7 @@ def test_print_pretty_onnx(self):
         )
         self.print_onnx(proto)
         self.print_model(proto)
-        self.dump_onnx("test_print_pretty_onnx", proto)
+        self.dump_onnx("test_print_pretty.onnx", proto)
         self.check_ort(proto)
         self.assertNotEmpty(proto)
         self.assertEmpty(None)
@@ -203,6 +203,101 @@ def test_max_diff_verbose(self):
         d = string_diff(diff)
         self.assertIsInstance(d, str)
 
+    def test_max_diff_hist_array(self):
+        x = np.arange(12).reshape((3, 4)).astype(dtype=np.float32)
+        y = x.copy()
+        y[0, 1] += 0.1
+        y[0, 2] += 0.01
+        y[0, 3] += 0.001
+        y[1, 1] += 0.0001
+        y[1, 2] += 1
+        y[2, 2] += 10
+        y[1, 3] += 100
+        y[2, 1] += 1000
+        diff = max_diff(x, y, hist=True)
+        self.assertEqual(
+            diff["rep"],
+            {
+                ">0.0": 8,
+                ">0.0001": 8,
+                ">0.001": 6,
+                ">0.01": 5,
+                ">0.1": 5,
+                ">1.0": 3,
+                ">10.0": 2,
+                ">100.0": 1,
+            },
+        )
+
+    def test_max_diff_hist_array_string_diff(self):
+        x = np.arange(12).reshape((3, 4)).astype(dtype=np.float32)
+        y = x.copy()
+        y[0, 1] += 0.1
+        y[0, 2] += 0.01
+        y[0, 3] += 0.001
+        y[1, 1] += 0.0001
+        y[1, 2] += 1
+        y[2, 2] += 10
+        y[1, 3] += 100
+        y[2, 1] += 1000
+        diff = max_diff(x, y, hist=True)
+        s = string_diff(diff)
+        self.assertEndsWith(
+            "/#8>0.0-#8>0.0001-#6>0.001-#5>0.01-#5>0.1-#3>1.0-#2>10.0-#1>100.0", s
+        )
+
+    def test_max_diff_hist_tensor(self):
+        x = torch.arange(12).reshape((3, 4)).to(dtype=torch.float32)
+        y = x.clone()
+        y[0, 1] += 0.1
+        y[0, 2] += 0.01
+        y[0, 3] += 0.001
+        y[1, 1] += 0.0001
+        y[1, 2] += 1
+        y[2, 2] += 10
+        y[1, 3] += 100
+        y[2, 1] += 1000
+        diff = max_diff(x, y, hist=True)
+        self.assertEqual(
+            diff["rep"],
+            {
+                ">0.0": 8,
+                ">0.0001": 8,
+                ">0.001": 6,
+                ">0.01": 5,
+                ">0.1": 5,
+                ">1.0": 3,
+                ">10.0": 2,
+                ">100.0": 1,
+            },
+        )
+
+    def test_max_diff_hist_tensor_composed(self):
+        x = torch.arange(12).reshape((3, 4)).to(dtype=torch.float32)
+        y = x.clone()
+        y[0, 1] += 0.1
+        y[0, 2] += 0.01
+        y[0, 3] += 0.001
+        y[1, 1] += 0.0001
+        y[1, 2] += 1
+        y[2, 2] += 10
+        y[1, 3] += 100
+        y[2, 1] += 1000
+        diff = max_diff([x, (x, {"e": x})], [y, (y, {"e": y})], hist=True)
+        self.assertEqual(
+            diff["rep"],
+            {
+                ">0.0": 24,
+                ">0.0001": 24,
+                ">0.001": 18,
+                ">0.01": 15,
+                ">0.1": 15,
+                ">1.0": 9,
+                ">10.0": 6,
+                ">100.0": 3,
+            },
+        )
+
     def test_type_info(self):
         for tt in [
             onnx.TensorProto.FLOAT,

_unittests/ut_helpers/test_torch_test_helper.py

@@ -9,6 +9,7 @@
     dummy_llm,
     to_numpy,
     is_torchdynamo_exporting,
+    model_statistics,
     steal_forward,
     replace_string_by_dynamic,
     to_any,
@@ -172,14 +173,15 @@ def forward(self, x, y):
             else:
                 print("output", k, v)
         print(string_type(restored, with_shape=True))
+        l1, l2 = 151, 160
         self.assertEqual(
             [
-                ("-Model-159", 0, "I"),
-                ("-Model-159", 0, "O"),
-                ("s1-SubModel-150", 0, "I"),
-                ("s1-SubModel-150", 0, "O"),
-                ("s2-SubModel-150", 0, "I"),
-                ("s2-SubModel-150", 0, "O"),
+                (f"-Model-{l2}", 0, "I"),
+                (f"-Model-{l2}", 0, "O"),
+                (f"s1-SubModel-{l1}", 0, "I"),
+                (f"s1-SubModel-{l1}", 0, "O"),
+                (f"s2-SubModel-{l1}", 0, "I"),
+                (f"s2-SubModel-{l1}", 0, "O"),
             ],
             sorted(restored),
         )
@@ -279,6 +281,32 @@ def test_torch_deepcopy_sliding_windon_cache(self):
     def test_torch_deepcopy_none(self):
         self.assertEmpty(torch_deepcopy(None))
 
+    def test_model_statistics(self):
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.p1 = torch.nn.Parameter(torch.tensor([1], dtype=torch.float32))
+                self.b1 = torch.nn.Buffer(torch.tensor([1], dtype=torch.float32))
+
+            def forward(self, x, y=None):
+                return x + y + self.p1 + self.b1
+
+        model = Model()
+        x, y = torch.rand((3, 4)), torch.rand((3, 4))
+        model(x, y)
+        stat = model_statistics(model)
+        self.assertEqual(
+            {
+                "type": "Model",
+                "n_modules": 1,
+                "param_size": 4,
+                "buffer_size": 4,
+                "float32": 8,
+                "size_mb": 0,
+            },
+            stat,
+        )
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

onnx_diagnostic/ext_test_case.py

@@ -1070,6 +1070,11 @@ def assertStartsWith(self, prefix: str, full: str):
         if not full.startswith(prefix):
             raise AssertionError(f"prefix={prefix!r} does not start string  {full!r}.")
 
+    def assertEndsWith(self, suffix: str, full: str):
+        """In the name"""
+        if not full.endswith(suffix):
+            raise AssertionError(f"suffix={suffix!r} does not end string  {full!r}.")
+
     def capture(self, fct: Callable):
         """
         Runs a function and capture standard output and error.