fix

Author: xadupre

_unittests/ut_xrun_doc/test_command_lines.py

@@ -6,6 +6,7 @@
     get_main_parser,
     get_parser_agg,
     get_parser_config,
+    get_parser_dot,
     get_parser_find,
     get_parser_lighten,
     get_parser_print,
@@ -23,6 +24,7 @@ def test_main_parser(self):
             get_main_parser().print_help()
         text = st.getvalue()
         self.assertIn("lighten", text)
+        self.assertIn("dot", text)
 
     def test_parser_lighten(self):
         st = StringIO()
@@ -87,6 +89,13 @@ def test_parser_sbs(self):
         text = st.getvalue()
         self.assertIn("--onnx", text)
 
+    def test_parser_dot(self):
+        st = StringIO()
+        with redirect_stdout(st):
+            get_parser_dot().print_help()
+        text = st.getvalue()
+        self.assertIn("--run", text)
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

_unittests/ut_xrun_doc/test_command_lines_exe.py

@@ -162,6 +162,45 @@ def forward(self, x):
         sdf = df[(df.ep_target == "placeholder") & (df.onnx_op_type == "initializer")]
         self.assertEqual(sdf.shape[0], 4)
 
+    @ignore_warnings(UserWarning)
+    @requires_transformers("4.53")
+    def test_i_parser_dot(self):
+        import torch
+
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super(Model, self).__init__()
+                self.fc1 = torch.nn.Linear(10, 32)  # input size 10 → hidden size 32
+                self.relu = torch.nn.ReLU()
+                self.fc2 = torch.nn.Linear(32, 1)  # hidden → output
+
+            def forward(self, x):
+                x = self.relu(self.fc1(x))
+                x = self.fc2(x)
+                return x
+
+        inputs = dict(x=torch.randn((5, 10)))
+        ds = dict(x={0: "batch"})
+        onnx_file = self.get_dump_file("test_i_parser_dot.model.onnx")
+        to_onnx(
+            Model(),
+            kwargs=inputs,
+            dynamic_shapes=ds,
+            exporter="custom",
+            filename=onnx_file,
+        )
+
+        output = self.get_dump_file("test_i_parser_dot.dot")
+        args = ["dot", onnx_file, "-v", "1", "-o", output]
+        if not self.unit_test_going():
+            args.extend(["--run", "svg"])
+
+        st = StringIO()
+        with redirect_stdout(st):
+            main(args)
+        text = st.getvalue()
+        print(text)
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

onnx_diagnostic/_command_lines_parser.py

@@ -11,6 +11,77 @@
 from argparse import ArgumentParser, RawTextHelpFormatter, BooleanOptionalAction
 
 
+def get_parser_dot() -> ArgumentParser:
+    parser = ArgumentParser(
+        prog="dot",
+        description=textwrap.dedent(
+            """
+            Converts a model into a dot file dot can draw into a graph.
+            """
+        ),
+    )
+    parser.add_argument("input", type=str, help="onnx model to lighten")
+    parser.add_argument(
+        "-o",
+        "--output",
+        default="",
+        type=str,
+        required=False,
+        help="dot model to output or empty to print out the result",
+    )
+    parser.add_argument(
+        "-v",
+        "--verbose",
+        type=int,
+        default=0,
+        required=False,
+        help="verbosity",
+    )
+    parser.add_argument(
+        "-r",
+        "--run",
+        default="",
+        required=False,
+        help="run dot, in that case, format must be given (svg, png)",
+    )
+    return parser
+
+
+def _cmd_dot(argv: List[Any]):
+    import subprocess
+    from .helpers.dot_helper import to_dot
+
+    parser = get_parser_dot()
+    args = parser.parse_args(argv[1:])
+    if args.verbose:
+        print(f"-- loads {args.input!r}")
+    onx = onnx.load(args.input, load_external_data=False)
+    if args.verbose:
+        print("-- converts into dot")
+    dot = to_dot(onx)
+    if args.output:
+        if args.verbose:
+            print(f"-- saves into {args.output}")
+        with open(args.output, "w") as f:
+            f.write(dot)
+    else:
+        print(dot)
+    if args.run:
+        assert args.output, "Cannot run dot without an output file."
+        cmds = ["dot", f"-T{args.run}", args.output, "-o", f"{args.output}.{args.run}"]
+        if args.verbose:
+            print(f"-- run {' '.join(cmds)}")
+        p = subprocess.Popen(cmds, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+        res = p.communicate()
+        out, err = res
+        if out:
+            print("--")
+            print(out)
+        if err:
+            print("--")
+            print(err)
+
+
 def get_parser_lighten() -> ArgumentParser:
     parser = ArgumentParser(
         prog="lighten",
@@ -1412,6 +1483,7 @@ def get_main_parser() -> ArgumentParser:
 
             agg          - aggregates statistics from multiple files
             config       - prints a configuration for a model id
+            dot          - converts an onnx model into dot format
             exportsample - produces a code to export a model
             find         - find node consuming or producing a result
             lighten      - makes an onnx model lighter by removing the weights,
@@ -1428,6 +1500,7 @@ def get_main_parser() -> ArgumentParser:
         choices=[
             "agg",
             "config",
+            "dot",
             "exportsample",
             "find",
             "lighten",
@@ -1446,6 +1519,7 @@ def main(argv: Optional[List[Any]] = None):
     fcts = dict(
         agg=_cmd_agg,
         config=_cmd_config,
+        dot=_cmd_dot,
         exportsample=_cmd_export_sample,
         find=_cmd_find,
         lighten=_cmd_lighten,
@@ -1470,6 +1544,7 @@ def main(argv: Optional[List[Any]] = None):
             parsers = dict(
                 agg=get_parser_agg,
                 config=get_parser_config,
+                dot=get_parser_dot,
                 exportsample=lambda: get_parser_validate("exportsample"),  # type: ignore[operator]
                 find=get_parser_find,
                 lighten=get_parser_lighten,

onnx_diagnostic/helpers/dot_helper.py

@@ -153,15 +153,12 @@ def _mkn(obj: object) -> int:
         shape = tuple(init.dims)
         if len(shape) == 0 or (len(shape) == 1 and shape[0] < 10):
             a = onh.to_array(init)
-            vals = f" = {a}" if len(shape) == 0 else f"\\n=[{', '.join([str(i) for i in a])}]"
             tiny_inits[init.name] = (
                 str(a) if len(shape) == 0 else f"[{', '.join([str(i) for i in a])}]"
             )
         else:
             ls = f"{onnx_dtype_name(init.data_type)}({', '.join(map(str,shape))})"
-            rows.append(
-                f'  i_{_mkn(init)} [label="{init.name}\\n{ls}{vals}", fillcolor="#cccc00"];'
-            )
+            rows.append(f'  i_{_mkn(init)} [label="{init.name}\\n{ls}", fillcolor="#cccc00"];')
             name_to_ids[init.name] = f"i_{_mkn(init)}"
             edge_label[init.name] = ls
     for node in nodes: