onnx-diagnostic 0.2.1__py3-none-any.whl → 0.3.0__py3-none-any.whl

onnx_diagnostic/__init__.py

@@ -3,5 +3,5 @@ Investigates onnx models.
 Functions, classes to dig into a model when this one is right, slow, wrong...
 """
 
-__version__ = "0.2.1"
+__version__ = "0.3.0"
 __author__ = "Xavier Dupré"

onnx_diagnostic/__main__.py

@@ -0,0 +1,4 @@
+from ._command_lines_parser import main
+
+if __name__ == "__main__":
+    main()

onnx_diagnostic/_command_lines_parser.py

@@ -0,0 +1,411 @@
+import json
+import sys
+import textwrap
+import onnx
+from typing import Any, List, Optional
+from argparse import ArgumentParser, RawTextHelpFormatter, BooleanOptionalAction
+from textwrap import dedent
+
+
+def get_parser_lighten() -> ArgumentParser:
+    parser = ArgumentParser(
+        prog="lighten",
+        description=dedent(
+            """
+        Removes the weights from a heavy model, stores statistics to restore
+        random weights.
+        """
+        ),
+        epilog="This is mostly used to write unit tests without adding "
+        "a big onnx file to the repository.",
+    )
+    parser.add_argument(
+        "-i",
+        "--input",
+        type=str,
+        required=True,
+        help="onnx model to lighten",
+    )
+    parser.add_argument(
+        "-o",
+        "--output",
+        type=str,
+        required=True,
+        help="onnx model to output",
+    )
+    parser.add_argument(
+        "-v",
+        "--verbose",
+        default=0,
+        required=False,
+        help="verbosity",
+    )
+    return parser
+
+
+def _cmd_lighten(argv: List[Any]):
+    from .helpers.onnx_helper import onnx_lighten
+
+    parser = get_parser_lighten()
+    args = parser.parse_args(argv[1:])
+    onx = onnx.load(args.input)
+    new_onx, stats = onnx_lighten(onx, verbose=args.verbose)
+    jstats = json.dumps(stats)
+    if args.verbose:
+        print("save file {args.input!r}")
+    if args.verbose:
+        print("write file {args.output!r}")
+    with open(args.output, "wb") as f:
+        f.write(new_onx.SerializeToString())
+    name = f"{args.output}.stats"
+    with open(name, "w") as f:
+        f.write(jstats)
+    if args.verbose:
+        print("done")
+
+
+def get_parser_unlighten() -> ArgumentParser:
+    parser = ArgumentParser(
+        prog="unlighten",
+        description=dedent(
+            """
+        Restores random weights for a model reduces with command lighten,
+        the command expects to find a file nearby with extension '.stats'.
+        """
+        ),
+        epilog="This is mostly used to write unit tests without adding "
+        "a big onnx file to the repository.",
+    )
+    parser.add_argument(
+        "-i",
+        "--input",
+        type=str,
+        required=True,
+        help="onnx model to unlighten",
+    )
+    parser.add_argument(
+        "-o",
+        "--output",
+        type=str,
+        required=True,
+        help="onnx model to output",
+    )
+    parser.add_argument(
+        "-v",
+        "--verbose",
+        default=0,
+        required=False,
+        help="verbosity",
+    )
+    return parser
+
+
+def _cmd_unlighten(argv: List[Any]):
+    from .helpers.onnx_helper import onnx_unlighten
+
+    parser = get_parser_lighten()
+    args = parser.parse_args(argv[1:])
+    new_onx = onnx_unlighten(args.input, verbose=args.verbose)
+    if args.verbose:
+        print(f"save file {args.output}")
+    with open(args.output, "wb") as f:
+        f.write(new_onx.SerializeToString())
+    if args.verbose:
+        print("done")
+
+
+def get_parser_print() -> ArgumentParser:
+    parser = ArgumentParser(
+        prog="print",
+        description=dedent(
+            """
+        Prints the model on the standard output.
+        """
+        ),
+        epilog="To show a model.",
+    )
+    parser.add_argument(
+        "fmt", choices=["pretty", "raw"], help="Format to use.", default="pretty"
+    )
+    parser.add_argument("input", type=str, help="onnx model to load")
+    return parser
+
+
+def _cmd_print(argv: List[Any]):
+    parser = get_parser_print()
+    args = parser.parse_args(argv[1:])
+    onx = onnx.load(args.input)
+    if args.fmt == "raw":
+        print(onx)
+    elif args.fmt == "pretty":
+        from .helpers.onnx_helper import pretty_onnx
+
+        print(pretty_onnx(onx))
+    else:
+        raise ValueError(f"Unexpected value fmt={args.fmt!r}")
+
+
+def get_parser_find() -> ArgumentParser:
+    parser = ArgumentParser(
+        prog="find",
+        description=dedent(
+            """
+        Look into a model and search for a set of names,
+        tells which node is consuming or producing it.
+        """
+        ),
+        epilog="Enables Some quick validation.",
+    )
+    parser.add_argument(
+        "-i",
+        "--input",
+        type=str,
+        required=True,
+        help="onnx model to unlighten",
+    )
+    parser.add_argument(
+        "-n",
+        "--names",
+        type=str,
+        required=False,
+        help="names to look at comma separated values",
+    )
+    parser.add_argument(
+        "-v",
+        "--verbose",
+        default=0,
+        required=False,
+        help="verbosity",
+    )
+    return parser
+
+
+def _cmd_find(argv: List[Any]):
+    from .helpers.onnx_helper import onnx_find
+
+    parser = get_parser_find()
+    args = parser.parse_args(argv[1:])
+    onnx_find(args.input, verbose=args.verbose, watch=set(args.names.split(",")))
+
+
+def get_parser_config() -> ArgumentParser:
+    parser = ArgumentParser(
+        prog="config",
+        description=dedent(
+            """
+        Prints out a configuration for a model id,
+        prints the associated task as well.
+        """
+        ),
+        epilog="",
+    )
+    parser.add_argument(
+        "-m",
+        "--mid",
+        type=str,
+        required=True,
+        help="model id, usually <author>/<name>",
+    )
+    parser.add_argument(
+        "-t",
+        "--task",
+        default=False,
+        action=BooleanOptionalAction,
+        help="displays the task as well",
+    )
+    return parser
+
+
+def _cmd_config(argv: List[Any]):
+    from .torch_models.hghub.hub_api import get_pretrained_config, task_from_id
+
+    parser = get_parser_config()
+    args = parser.parse_args(argv[1:])
+    print(get_pretrained_config(args.mid))
+    if args.task:
+        print("------")
+        print(f"task: {task_from_id(args.mid)}")
+
+
+def get_parser_validate() -> ArgumentParser:
+    parser = ArgumentParser(
+        prog="test",
+        description=dedent(
+            """
+        Prints out dummy inputs for a particular task or a model id.
+        If both mid and task are empty, the command line displays the list
+        of supported tasks.
+        """
+        ),
+        epilog="If the model id is specified, one untrained version of it is instantiated.",
+    )
+    parser.add_argument("-m", "--mid", type=str, help="model id, usually <author>/<name>")
+    parser.add_argument("-t", "--task", default=None, help="force the task to use")
+    parser.add_argument("-e", "--export", help="export the model with this exporter")
+    parser.add_argument("--opt", help="optimization to apply after the export")
+    parser.add_argument(
+        "-r",
+        "--run",
+        default=False,
+        action=BooleanOptionalAction,
+        help="runs the model to check it runs",
+    )
+    parser.add_argument(
+        "-q",
+        "--quiet",
+        default=False,
+        action=BooleanOptionalAction,
+        help="catches exception, report them in the summary",
+    )
+    parser.add_argument(
+        "-p",
+        "--patch",
+        default=True,
+        action=BooleanOptionalAction,
+        help="applies patches before exporting",
+    )
+    parser.add_argument(
+        "--stop-if-static",
+        default=0,
+        type=int,
+        help="raises an exception if a dynamic dimension becomes static",
+    )
+    parser.add_argument(
+        "--trained",
+        default=False,
+        action=BooleanOptionalAction,
+        help="validate the trained model (requires downloading)",
+    )
+    parser.add_argument(
+        "-o",
+        "--dump-folder",
+        help="if not empty, a folder is created to dumps statistics, "
+        "exported program, onnx...",
+    )
+    parser.add_argument(
+        "--drop",
+        help="drops the following inputs names, it should be a list "
+        "with comma separated values",
+    )
+    parser.add_argument("-v", "--verbose", default=0, type=int, help="verbosity")
+    parser.add_argument("--dtype", help="changes dtype if necessary")
+    parser.add_argument("--device", help="changes the device if necessary")
+    return parser
+
+
+def _cmd_validate(argv: List[Any]):
+    from .helpers import string_type
+    from .torch_models.test_helper import get_inputs_for_task, validate_model, _ds_clean
+    from .torch_models.hghub.model_inputs import get_get_inputs_function_for_tasks
+
+    parser = get_parser_validate()
+    args = parser.parse_args(argv[1:])
+    if not args.task and not args.mid:
+        print("-- list of supported tasks:")
+        print("\n".join(sorted(get_get_inputs_function_for_tasks())))
+    elif not args.mid:
+        data = get_inputs_for_task(args.task)
+        if args.verbose:
+            print(f"task: {args.task}")
+        max_length = max(len(k) for k in data["inputs"]) + 1
+        print("-- inputs")
+        for k, v in data["inputs"].items():
+            print(f"  + {k.ljust(max_length)}: {string_type(v, with_shape=True)}")
+        print("-- dynamic_shapes")
+        for k, v in data["dynamic_shapes"].items():
+            print(f"  + {k.ljust(max_length)}: {_ds_clean(v)}")
+    else:
+        summary, _data = validate_model(
+            model_id=args.mid,
+            task=args.task,
+            do_run=args.run,
+            verbose=args.verbose,
+            quiet=args.quiet,
+            trained=args.trained,
+            dtype=args.dtype,
+            device=args.device,
+            patch=args.patch,
+            stop_if_static=args.stop_if_static,
+            optimization=args.opt,
+            exporter=args.export,
+            dump_folder=args.dump_folder,
+            drop_inputs=None if not args.drop else args.drop.split(","),
+        )
+        print("")
+        print("-- summary --")
+        for k, v in sorted(summary.items()):
+            print(f":{k},{v};")
+
+
+def get_main_parser() -> ArgumentParser:
+    parser = ArgumentParser(
+        prog="onnx_diagnostic",
+        description="onnx_diagnostic main command line.\n",
+        formatter_class=RawTextHelpFormatter,
+        epilog=textwrap.dedent(
+            """
+        Type 'python -m onnx_diagnostic <cmd> --help'
+        to get help for a specific command.
+
+        config     - prints a configuration for a model id
+        find       - find node consuming or producing a result
+        lighten    - makes an onnx model lighter by removing the weights,
+        unlighten  - restores an onnx model produces by the previous experiment
+        print      - prints the model on standard output
+        validate   - validate a model
+        """
+        ),
+    )
+    parser.add_argument(
+        "cmd",
+        choices=["config", "find", "lighten", "print", "unlighten", "validate"],
+        help="Selects a command.",
+    )
+    return parser
+
+
+def main(argv: Optional[List[Any]] = None):
+    fcts = dict(
+        lighten=_cmd_lighten,
+        unlighten=_cmd_unlighten,
+        print=_cmd_print,
+        find=_cmd_find,
+        config=_cmd_config,
+        validate=_cmd_validate,
+    )
+
+    if argv is None:
+        argv = sys.argv[1:]
+    if (
+        len(argv) == 0
+        or (len(argv) <= 1 and argv[0] not in fcts)
+        or argv[-1] in ("--help", "-h")
+    ):
+        if len(argv) < 2:
+            parser = get_main_parser()
+            parser.parse_args(argv)
+        else:
+            parsers = dict(
+                lighten=get_parser_lighten,
+                unlighten=get_parser_unlighten,
+                print=get_parser_print,
+                find=get_parser_find,
+                config=get_parser_config,
+                validate=get_parser_validate,
+            )
+            cmd = argv[0]
+            if cmd not in parsers:
+                raise ValueError(
+                    f"Unknown command {cmd!r}, it should be in {list(sorted(parsers))}."
+                )
+            parser = parsers[cmd]()
+            parser.parse_args(argv[1:])
+        raise RuntimeError("The programme should have exited before.")
+
+    cmd = argv[0]
+    if cmd in fcts:
+        fcts[cmd](argv)
+    else:
+        raise ValueError(
+            f"Unknown command {cmd!r}, use --help to get the list of known command."
+        )

onnx_diagnostic/doc.py

@@ -0,0 +1,32 @@
+def reset_torch_transformers(gallery_conf, fname):
+    "Resets torch dynamo for :epkg:`sphinx-gallery`."
+    import matplotlib.pyplot as plt
+    import torch
+
+    plt.style.use("ggplot")
+    torch._dynamo.reset()
+
+
+def plot_legend(
+    text: str, text_bottom: str = "", color: str = "green", fontsize: int = 15
+) -> "matplotlib.axes.Axes":  # noqa: F821
+    import matplotlib.pyplot as plt
+
+    fig = plt.figure(figsize=(2, 2))
+    ax = fig.add_subplot()
+    ax.axis([0, 5, 0, 5])
+    ax.text(2.5, 4, "END", fontsize=10, horizontalalignment="center")
+    ax.text(
+        2.5,
+        2.5,
+        text,
+        fontsize=fontsize,
+        bbox={"facecolor": color, "alpha": 0.5, "pad": 10},
+        horizontalalignment="center",
+        verticalalignment="center",
+    )
+    if text_bottom:
+        ax.text(4.5, 0.5, text_bottom, fontsize=7, horizontalalignment="right")
+    ax.grid(False)
+    ax.set_axis_off()
+    return ax

onnx_diagnostic/export/__init__.py

@@ -1 +1 @@
-from .dynamic_shapes import ModelInputs
+from .dynamic_shapes import CoupleInputsDynamicShapes, ModelInputs