extend code coverage

xadupre · xadupre · commit 8b22076978df · 2025-03-21T12:52:52.000+01:00
diff --git a/_unittests/ut_xrun_doc/test_helpers.py b/_unittests/ut_xrun_doc/test_helpers.py
@@ -1,7 +1,12 @@
 import unittest
 import numpy as np
+import onnx
+import onnx.helper as oh
+import torch
 from onnx_diagnostic.ext_test_case import ExtTestCase, skipif_ci_windows
-from onnx_diagnostic.helpers import string_type, string_sig
+from onnx_diagnostic.helpers import string_type, string_sig, pretty_onnx, get_onnx_signature
+
+TFLOAT = onnx.TensorProto.FLOAT
 
 
 class TestHelpers(ExtTestCase):
@@ -19,8 +24,6 @@ def test_string_dict(self):
         self.assertEqual(s, "dict(a:A1r1,b:dict(r:float),c:{int})")
 
     def test_string_type_array(self):
-        import torch
-
         a = np.array([1], dtype=np.float32)
         t = torch.tensor([1])
         obj = {"a": a, "b": t}
@@ -30,22 +33,64 @@ def test_string_type_array(self):
         self.assertEqual(s, "dict(a:A1s1,b:T7s1)")
 
     def test_string_sig_f(self):
-
         def f(a, b=3, c=4, e=5):
             pass
 
         ssig = string_sig(f, {"a": 1, "c": 8, "b": 3})
         self.assertEqual(ssig, "f(a=1, c=8)")
 
     def test_string_sig_cls(self):
-
         class A:
             def __init__(self, a, b=3, c=4, e=5):
                 self.a, self.b, self.c, self.e = a, b, c, e
 
         ssig = string_sig(A(1, c=8))
         self.assertEqual(ssig, "A(a=1, c=8)")
 
+    def test_pretty_onnx(self):
+        proto = oh.make_model(
+            oh.make_graph(
+                [
+                    oh.make_node("Sigmoid", ["Y"], ["sy"]),
+                    oh.make_node("Mul", ["Y", "sy"], ["ysy"]),
+                    oh.make_node("Mul", ["X", "ysy"], ["final"]),
+                ],
+                "nd",
+                [
+                    oh.make_tensor_value_info("X", TFLOAT, [1, "b", "c"]),
+                    oh.make_tensor_value_info("Y", TFLOAT, ["a", "b", "c"]),
+                ],
+                [oh.make_tensor_value_info("final", TFLOAT, ["a", "b", "c"])],
+            ),
+            opset_imports=[oh.make_opsetid("", 18)],
+            ir_version=9,
+        )
+        pretty_onnx(proto, shape_inference=True)
+        pretty_onnx(proto.graph.input[0])
+        pretty_onnx(proto.graph)
+        pretty_onnx(proto.graph.node[0])
+
+    def test_get_onnx_signature(self):
+        proto = oh.make_model(
+            oh.make_graph(
+                [
+                    oh.make_node("Sigmoid", ["Y"], ["sy"]),
+                    oh.make_node("Mul", ["Y", "sy"], ["ysy"]),
+                    oh.make_node("Mul", ["X", "ysy"], ["final"]),
+                ],
+                "nd",
+                [
+                    oh.make_tensor_value_info("X", TFLOAT, [1, "b", "c"]),
+                    oh.make_tensor_value_info("Y", TFLOAT, ["a", "b", "c"]),
+                ],
+                [oh.make_tensor_value_info("final", TFLOAT, ["a", "b", "c"])],
+            ),
+            opset_imports=[oh.make_opsetid("", 18)],
+            ir_version=9,
+        )
+        sig = get_onnx_signature(proto)
+        self.assertEqual(sig, (("X", 1, (1, "b", "c")), ("Y", 1, ("a", "b", "c"))))
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)
diff --git a/_unittests/ut_xrun_doc/test_onnx_tools.py b/_unittests/ut_xrun_doc/test_onnx_tools.py
@@ -0,0 +1,72 @@
+import unittest
+import numpy as np
+import onnx.helper as oh
+import onnx.numpy_helper as onh
+from onnx import TensorProto
+from onnx.checker import check_model
+from onnx_diagnostic.ext_test_case import ExtTestCase
+from onnx_diagnostic.onnx_tools import onnx_lighten, onnx_unlighten, onnx_find
+from onnx_diagnostic.torch_test_helper import check_model_ort
+
+TFLOAT = TensorProto.FLOAT
+
+
+class TestOnnxTools(ExtTestCase):
+
+    def _get_model(self):
+        model = oh.make_model(
+            oh.make_graph(
+                [
+                    oh.make_node("Unsqueeze", ["X", "zero"], ["xu1"]),
+                    oh.make_node("Unsqueeze", ["xu1", "un"], ["xu2"]),
+                    oh.make_node("Reshape", ["xu2", "shape1"], ["xm1"]),
+                    oh.make_node("Reshape", ["Y", "shape2"], ["xm2c"]),
+                    oh.make_node("Cast", ["xm2c"], ["xm2"], to=1),
+                    oh.make_node("MatMul", ["xm1", "xm2"], ["xm"]),
+                    oh.make_node("Reshape", ["xm", "shape3"], ["Z"]),
+                ],
+                "dummy",
+                [oh.make_tensor_value_info("X", TFLOAT, [320, 1280])],
+                [oh.make_tensor_value_info("Z", TFLOAT, [3, 5, 320, 640])],
+                [
+                    onh.from_array(
+                        np.random.rand(3, 5, 1280, 640).astype(np.float32), name="Y"
+                    ),
+                    onh.from_array(np.array([0], dtype=np.int64), name="zero"),
+                    onh.from_array(np.array([1], dtype=np.int64), name="un"),
+                    onh.from_array(np.array([1, 320, 1280], dtype=np.int64), name="shape1"),
+                    onh.from_array(np.array([15, 1280, 640], dtype=np.int64), name="shape2"),
+                    onh.from_array(np.array([3, 5, 320, 640], dtype=np.int64), name="shape3"),
+                ],
+            ),
+            opset_imports=[oh.make_opsetid("", 18)],
+            ir_version=9,
+        )
+        return model
+
+    def test_un_lighten_model(self):
+        model = self._get_model()
+        check_model(model)
+        size1 = len(model.SerializeToString())
+        (onx, stats), out, _ = self.capture(lambda: onnx_lighten(model, verbose=1))
+        self.assertIsInstance(stats, dict)
+        self.assertEqual(len(stats), 1)
+        self.assertIsInstance(stats["Y"], dict)
+        self.assertIn("remove initializer", out)
+        # check_model(onx)
+        new_model = onnx_unlighten(onx, stats)
+        check_model(new_model)
+        size2 = len(new_model.SerializeToString())
+        self.assertEqual(size1, size2)
+        check_model_ort(model)
+
+    def test_onnx_find(self):
+        model = self._get_model()
+        res = onnx_find(model, watch={"xm2"})
+        self.assertEqual(len(res), 2)
+        self.assertIn("xm2", res[0].output)
+        self.assertIn("xm2", res[1].input)
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)
diff --git a/onnx_diagnostic/helpers.py b/onnx_diagnostic/helpers.py
@@ -6,6 +6,7 @@
 from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
 import numpy as np
 import numpy.typing as npt
+import onnx
 from onnx import (
     AttributeProto,
     FunctionProto,
@@ -528,7 +529,7 @@ def pretty_onnx(
     assert onx is not None, "onx cannot be None"
 
     if shape_inference:
-        onx = onx.shape_inference.infer_shapes(onx)
+        onx = onnx.shape_inference.infer_shapes(onx)
 
     if isinstance(onx, ValueInfoProto):
         name = onx.name
