onnx-diagnostic 0.8.5__py3-none-any.whl → 0.8.7__py3-none-any.whl

onnx_diagnostic/__init__.py

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

onnx_diagnostic/_command_lines_parser.py

@@ -1508,6 +1508,151 @@ def _cmd_sbs(argv: List[Any]):
     print("-- done")
 
 
+def get_parser_compare() -> ArgumentParser:
+    parser = ArgumentParser(
+        prog="compare",
+        description=textwrap.dedent(
+            """
+            Compares two onnx models by aligning the nodes between both models.
+            This is done through an edit distance.
+            """
+        ),
+        epilog=textwrap.dedent(
+            """
+            Each element (initializer, input, node, output) of the model
+            is converted into an observation. Then it defines a distance between
+            two elements. And finally, it finds the best alignment with
+            an edit distance.
+            """
+        ),
+    )
+    parser.add_argument("model1", type=str, help="first model to compare")
+    parser.add_argument("model2", type=str, help="second model to compare")
+    return parser
+
+
+def _cmd_compare(argv: List[Any]):
+    import onnx
+    from .torch_onnx.compare import ObsCompare, ObsComparePair
+
+    parser = get_parser_compare()
+    args = parser.parse_args(argv[1:])
+    print(f"-- loading {args.model1!r}")
+    seq1 = ObsCompare.obs_sequence_from_model(onnx.load(args.model1, load_external_data=False))
+    print(f"-- loading {args.model2!r}")
+    seq2 = ObsCompare.obs_sequence_from_model(onnx.load(args.model2, load_external_data=False))
+    print("-- starts comparison")
+    dist, _path, pair_cmp = ObsComparePair.distance_sequence(seq1, seq2)
+    print(f"-- done with distance {dist}")
+    print(ObsComparePair.to_str(pair_cmp))
+
+
+def get_parser_optimize() -> ArgumentParser:
+    parser = ArgumentParser(
+        prog="optimize",
+        formatter_class=RawTextHelpFormatter,
+        description=textwrap.dedent(
+            """
+            Optimizes an onnx model by fusing nodes. It looks for patterns in the graphs
+            and replaces them by the corresponding nodes. It also does basic optimization
+            such as removing identity nodes or unused nodes.
+            """
+        ),
+        epilog=textwrap.dedent(
+            """
+            The goal is to make the model faster.
+            Argument patterns defines the patterns to apply or the set of patterns.
+            It is possible to show statistics or to remove a particular pattern.
+            Here are some environment variables which can be used to trigger
+            these displays.
+
+            Available options algorithms, default and default+runtime:
+
+            - DROPPATTERN=<pattern1,patterns2,...>: do not apply
+              those patterns when optimizing a model
+            - DUMPPATTERNS=<folder>: dumps all matched and applied
+              nodes when a pattern is applied
+            - PATTERN=<pattern1,pattern2,...>: increase verbosity for specific
+                patterns to understand why one pattern was not applied,
+                this shows which line is rejecting a pattern if it seems one pattern was missed
+            """
+        ),
+    )
+    parser.add_argument(
+        "algorithm",
+        choices=["ir", "os_ort", "slim", "default", "default+onnxruntime"],
+        help="algorithm or patterns optimization to apply",
+    )
+    parser.add_argument("input", type=str, help="onnx model to optimize")
+    parser.add_argument(
+        "-o",
+        "--output",
+        type=str,
+        required=False,
+        help="onnx model to output, if empty, if adds .opt-{algorithm}.onnx to the name",
+    )
+    parser.add_argument(
+        "-v",
+        "--verbose",
+        default=0,
+        required=False,
+        type=int,
+        help="verbosity",
+    )
+    parser.add_argument(
+        "--infer-shapes",
+        default=True,
+        action=BooleanOptionalAction,
+        help="infer shapes before optimizing the model",
+    )
+    parser.add_argument(
+        "--processor",
+        default="",
+        help=textwrap.dedent(
+            """
+            optimization for a specific processor, CPU, CUDA or both CPU,CUDA,
+            some operators are only available in one processor, it might be not used
+            with all
+            """
+        ).strip("\n"),
+    )
+    parser.add_argument(
+        "--remove-shape-info",
+        default=True,
+        action=BooleanOptionalAction,
+        help="remove shape information before outputting the model",
+    )
+    return parser
+
+
+def _cmd_optimize(argv: List[Any]):
+    parser = get_parser_optimize()
+    args = parser.parse_args(argv[1:])
+
+    from .helpers.optim_helper import optimize_model
+
+    output = (
+        args.output
+        if args.output
+        else f"{os.path.splitext(args.input)[0]}.o-{args.algorithm}.onnx"
+    )
+
+    optimize_model(
+        args.algorithm,
+        args.input,
+        output=output,
+        verbose=args.verbose,
+        processor=args.processor,
+        infer_shapes=args.infer_shapes,
+        remove_shape_info=args.remove_shape_info,
+    )
+
+
+#############
+# main parser
+#############
+
+
 def get_main_parser() -> ArgumentParser:
     parser = ArgumentParser(
         prog="onnx_diagnostic",
@@ -1519,16 +1664,17 @@ def get_main_parser() -> ArgumentParser:
             to get help for a specific command.
 
             agg          - aggregates statistics from multiple files
-            config       - prints a configuration for a model id
+            config       - prints a configuration for a model id (on HuggingFace Hub)
             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,
+            lighten      - makes an onnx model lighter by removing the weights
+            optimize     - optimizes an onnx model
             print        - prints the model on standard output
             sbs          - compares an exported program and a onnx model
             stats        - produces statistics on a model
             unlighten    - restores an onnx model produces by the previous experiment
-            validate     - validate a model
+            validate     - validate a model (knowing its model id on HuggginFace Hub)
             """
         ),
     )
@@ -1541,6 +1687,7 @@ def get_main_parser() -> ArgumentParser:
             "exportsample",
             "find",
             "lighten",
+            "optimize",
             "print",
             "sbs",
             "stats",
@@ -1555,11 +1702,13 @@ def get_main_parser() -> ArgumentParser:
 def main(argv: Optional[List[Any]] = None):
     fcts = dict(
         agg=_cmd_agg,
+        compare=_cmd_compare,
         config=_cmd_config,
         dot=_cmd_dot,
         exportsample=_cmd_export_sample,
         find=_cmd_find,
         lighten=_cmd_lighten,
+        optimize=_cmd_optimize,
         print=_cmd_print,
         sbs=_cmd_sbs,
         stats=_cmd_stats,
@@ -1580,11 +1729,13 @@ def main(argv: Optional[List[Any]] = None):
         else:
             parsers = dict(
                 agg=get_parser_agg,
+                compare=get_parser_compare,
                 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,
+                optimize=get_parser_optimize,
                 print=get_parser_print,
                 sbs=get_parser_sbs,
                 stats=get_parser_stats,

onnx_diagnostic/ci_models/ci_helpers.py

@@ -0,0 +1,435 @@
+import datetime
+import os
+import time
+import subprocess
+from argparse import ArgumentParser, BooleanOptionalAction, RawTextHelpFormatter
+from typing import Any, Dict, List, Tuple
+import onnx
+
+
+def get_versions():
+    """
+    Returns the version of the package currently used.
+    The output is a dictionary.
+    The function uses delayed import to make to fail fast at startup.
+    """
+    import onnx
+    import onnx_diagnostic
+    import onnxruntime
+    import torch
+    import transformers
+
+    return {
+        "transformers": transformers.__version__,
+        "onnxruntime": onnxruntime.__version__,
+        "onnx": onnx.__version__,
+        "onnx-diagnostic": onnx_diagnostic.__version__,
+        "torch": torch.__version__,
+    }
+
+
+def get_torch_dtype_from_command_line_args(dtype: str) -> "torch.dtype":  # noqa: F821
+    """
+    Returns the torch dtype base on the argument provided on the command line.
+
+    Imports are delayed to be faster when running the help of the command line.
+    """
+    import torch
+
+    torch_dtype = {
+        "float16": torch.float16,
+        "bfloat16": torch.bfloat16,
+        "float32": torch.float32,
+        "fp16": torch.float16,
+        "bf16": torch.bfloat16,
+        "fp32": torch.float32,
+    }
+    assert (
+        dtype in torch_dtype
+    ), f"Unexpected dtype {dtype!r}, not found in {set(torch_dtype)}."
+    return torch_dtype[dtype]
+
+
+def get_parser(name: str, epilog: str = "") -> ArgumentParser:
+    """Creates a default parser for many models."""
+    parser = ArgumentParser(
+        prog=name,
+        description=f"""Export command line for model {name!r}.""",
+        epilog=epilog,
+        formatter_class=RawTextHelpFormatter,
+    )
+    parser.add_argument(
+        "-m",
+        "--mid",
+        type=str,
+        default="Qwen/Qwen2.5-VL-7B-Instruct",
+        help="model id, default is Qwen/Qwen2.5-VL-7B-Instruct",
+    )
+    parser.add_argument("-d", "--device", default="cpu", help="Device, cpu (default) or cuda.")
+    parser.add_argument(
+        "-t", "--dtype", default="float32", help="dtype, float32 (default) or float16"
+    )
+    parser.add_argument(
+        "-e", "--exporter", default="onnx-dynamo", help="exporter, default is onnx-dynamo"
+    )
+    parser.add_argument(
+        "--pretrained",
+        default=True,
+        help="use pretrained model or a random model",
+        action=BooleanOptionalAction,
+    )
+    parser.add_argument(
+        "--second-input",
+        default=True,
+        help="check discrepancies with other inputs",
+        action=BooleanOptionalAction,
+    )
+    parser.add_argument(
+        "--zip",
+        default=False,
+        help="Creates a file .zip with onnx file and data file.",
+        action=BooleanOptionalAction,
+    )
+    parser.add_argument(
+        "-o",
+        "--output-folder",
+        default="dump_models",
+        help="Folders where to put the results.",
+        action=BooleanOptionalAction,
+    )
+    parser.add_argument(
+        "-x",
+        "--existing-onnx",
+        default="",
+        help="If an onnx file exists, only measures the discrepancies.",
+    )
+    parser.add_argument(
+        "-p",
+        "--part",
+        default="visual",
+        help="part of the model to export",
+    )
+    parser.add_argument(
+        "-a",
+        "--atol",
+        type=float,
+        default=2.0,
+        help="fails if the maximum discrepancy is above that threshold",
+    )
+    parser.add_argument(
+        "--mismatch01",
+        type=float,
+        default=0.1,
+        help="fails if the ratio of mismatches at level 0.1 is above that threshold",
+    )
+    parser.add_argument(
+        "--profile-exporter",
+        default=False,
+        help="Profiles the exporter and outputs an html document from pyinstrument",
+        action=BooleanOptionalAction,
+    )
+    return parser
+
+
+def remove_inplace_body_last_input_output_type_for_loop_because_they_might_be_sequences(
+    filename: str,
+):
+    """
+    Modified inplace an onnx file. It wipes out shapes provided
+    in ``model.graph.value_info`` because they are wrong when a Loop outputs
+    a sequence. It alose removes the types in attribute 'Body'
+    of an operator Loop because it may be a tensor when a sequence is expected.
+    This should not be needed in the future.
+    """
+    model = onnx.load(filename, load_external_data=False)
+    for node in model.graph.node:
+        if node.op_type == "Loop":
+            g = node.attribute[0].g
+            g.input[-1].type.CopyFrom(onnx.TypeProto())
+            g.output[-1].type.CopyFrom(onnx.TypeProto())
+    del model.graph.value_info[:]
+    model = onnx.shape_inference.infer_shapes(model)
+    onnx.save(model, filename, save_as_external_data=False)
+
+
+def simplify_model_id_for_a_filename(model_id: str) -> str:
+    """Changes a model id in a way it can be used in a filename."""
+    return model_id.lower().replace("/", ".")
+
+
+def compute_expected_outputs(
+    output_filename: str, model_to_export: "torch.nn.Module", input_filename: str  # noqa: F821
+) -> Tuple[Any, List[Any], List[float]]:
+    """
+    Computes the expected outputs for a model.
+    The function uses delayed import to make to fail fast at startup.
+
+    It caches the expected outputs in a file. They are restored if the file exists
+    or computed and saved if not.
+
+    Imports are delayed to be faster when running the help of the command line.
+    """
+    import tqdm
+    import torch
+    from ..helpers import string_type
+
+    inputs = torch.load(input_filename, weights_only=False)
+    export_inputs = inputs["export_inputs"]
+    other_inputs = inputs["other_inputs"]
+
+    if os.path.exists(output_filename):
+        print(f"-- restore expected outputs from {output_filename!r}")
+        expected = torch.load(output_filename, weights_only=False)
+        export_expected = expected["export_expected"]
+        other_expected = expected["other_expected"]
+        durations = expected["durations"]
+    else:
+        print(
+            f"-- compute with inputs: "
+            f"{string_type(export_inputs, with_shape=True, with_device=True)}"
+        )
+        export_expected = model_to_export(**export_inputs)
+        print(f"-- got: {string_type(export_expected, with_shape=True)}")
+        print(
+            f"-- compute with inputs: "
+            f"{string_type(other_inputs, with_shape=True, with_device=True)}"
+        )
+        other_expected = []
+        durations = []
+        for other in tqdm.tqdm(other_inputs):
+            begin = time.perf_counter()
+            expected = model_to_export(**other)
+            other_expected.append(expected)
+            durations.append(time.perf_counter() - begin)
+        print(f"-- got: {string_type(other_expected, with_shape=True, with_device=True)}")
+
+        expected = dict(
+            export_expected=export_expected,
+            other_expected=other_expected,
+            durations=durations,
+        )
+        print(f"-- dump expected outputs into {output_filename!r}")
+        torch.save(expected, output_filename)
+    print(f"-- computation took {sum(durations)}")
+    print(
+        f"-- export_expected={string_type(export_expected, with_shape=True, with_device=True)}"
+    )
+    print(
+        f"-- other_expected={string_type(other_expected, with_shape=True, with_device=True)}"
+    )
+    return export_expected, other_expected, durations
+
+
+def check_for_discrepancies_and_log_everything_into_a_json_file(
+    agg_stat_file: str,
+    stat_file: str,
+    export_duration: float,
+    device: str,
+    model_file: str,
+    cached_inputs: str,
+    cached_expected_outputs: str,
+    main_info: Dict[str, Any],
+    atol: float,
+    mismatch01: float,
+):
+    """
+    Checks discrepancies for a specific model.
+
+    Imports are delayed to be faster when running the help of the command line.
+
+    :param agg_stat_file: a file when the discrepancies are collected, this is used to
+        produce a table to make it easier to compare across types, devices, ...
+    :param stat_file: discrepancies results dumps into that file
+    :param export_duration: export duration
+    :param device: targeted device (to select onnxruntime provider)
+    :param model_file: onnx model file
+    :param cache_inputs: inputs saved with :func:`torch.save` and
+        restored with :func:`torch.load`,
+        needs to contains `export_inputs` (to check the model is valid),
+        and `other_inputs`, other sets of inputs to measure the discrepancies,
+        and speed up (rough estimation)
+    :param cached_expected_outputs: expected outputs saved with :func:`torch.save`
+        and restored with :func:`torch.load`,
+        needs to contains `export_expected` (to check the model is valid),
+        and `other_expected`, other sets of outputs to measure the discrepancies,
+        and speed up (rough estimation)
+    :param main_info: a dictionary with values used to tell which version, device, ...
+    :param atol: assert if tolerance is above this
+    :param mismatch01: assert if the ratio of mismatches is above that threshold
+    """
+    import tqdm
+    import onnxruntime
+    import torch
+    from ..helpers import flatten_object, max_diff, string_type, string_diff
+
+    cached = (
+        torch.load(cached_inputs, weights_only=False),
+        torch.load(cached_expected_outputs, weights_only=False),
+    )
+    durations = cached[0].get("durations", [])
+    export_inputs = cached[0]["export_inputs"]
+    other_inputs = cached[0]["other_inputs"]
+    export_expected = cached[1]["export_expected"]
+    other_expected = cached[1]["other_expected"]
+
+    onx = onnx.load(model_file, load_external_data=False)
+    opsets = [d for d in onx.opset_import if d.domain == ""]
+    assert (
+        opsets
+    ), f"Unable to find standard opset in file {model_file!r}, opsets={onx.opset_import}"
+    opset = opsets[0].version
+
+    with open(stat_file, "w") as f:
+
+        def fprint(s):
+            print(s)
+            f.write(f"{s}\n")
+
+        fprint(f"-- export duration: {export_duration}")
+        providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
+        if device == "cpu":
+            providers = providers[1:]
+        fprint(f"-- checking discrepancies with providers={providers!r}")
+        fprint(f"-- model_file={model_file!r}")
+        sess = onnxruntime.InferenceSession(model_file, providers=providers)
+
+        fprint(
+            f"-- export_inputs {string_type(export_inputs, with_shape=True, with_device=True)}"
+        )
+        fprint(
+            f"-- export_expected "
+            f"{string_type(export_expected, with_shape=True, with_device=True)}"
+        )
+        feeds = dict(
+            zip(
+                [i.name for i in sess.get_inputs()],
+                [
+                    v.detach().cpu().numpy()
+                    for v in flatten_object(export_inputs, drop_keys=True)
+                ],
+            )
+        )
+        small = sess.run(None, feeds)
+        flat_export_expected = flatten_object(export_expected, drop_keys=True)
+        diff = max_diff(flat_export_expected, small, hist=[0.1, 0.01])
+        fprint(f"-- discrepancies={diff}")
+        assert diff["abs"] <= atol and diff["rep"][">0.1"] / diff["n"] <= mismatch01, (
+            f"absolute error {diff['abs']} is above {atol} or number of "
+            f"mismatches ({diff['rep']['>0.1'] / diff['n']}) is above "
+            f"{mismatch01}, dicrepancies={string_diff(diff)}"
+        )
+
+        if other_inputs and other_expected:
+            feeds = [
+                dict(
+                    zip(
+                        [i.name for i in sess.get_inputs()],
+                        [
+                            v.detach().cpu().numpy()
+                            for v in flatten_object(inputs, drop_keys=True)
+                        ],
+                    )
+                )
+                for inputs in other_inputs
+            ]
+            fprint("")
+            fprint(f"-- inputs {string_type(feeds, with_shape=True, with_device=True)}")
+            fprint(
+                f"-- expected {string_type(other_expected, with_shape=True, with_device=True)}"
+            )
+            begin = time.perf_counter()
+            gots = []
+            for feed in tqdm.tqdm(feeds):
+                gots.append(sess.run(None, feed))
+            oduration = time.perf_counter() - begin
+            fprint(
+                f"-- torch duration={sum(durations[:len(gots)])}, onnx duration={oduration}, "
+                f"speedup={sum(durations[:len(gots)])/oduration} n={len(gots)}"
+            )
+
+            info = {
+                **main_info,
+                "timestamp": datetime.datetime.now().isoformat(),
+                "export_duration": export_duration,
+                "latency_torch": sum(durations[: len(gots)]),
+                "latency_ort": oduration,
+                "speedup": sum(durations[: len(gots)]) / oduration,
+                "latency_ort_n": len(gots),
+                "opset": opset,
+                **get_versions(),
+            }
+            with open(agg_stat_file, "a") as fs:
+                for fe, e, b in zip(feeds, other_expected, gots):
+                    flat_e = flatten_object(e, drop_keys=True)
+                    se = string_type(fe, with_shape=True)
+                    diff = max_diff(flat_e, b, hist=[0.1, 0.01])
+                    assert (
+                        diff["abs"] <= atol and diff["rep"][">0.1"] / diff["n"] <= mismatch01
+                    ), (
+                        f"absolute error {diff['abs']} is above {atol} or number "
+                        f" of mismatches ({diff['rep']['>0.1'] / diff['n']}) "
+                        f"is above {mismatch01}, dicrepancies={string_diff(diff)}"
+                    )
+                    js = string_diff(diff, js=True, ratio=True, inputs=se, **info)
+                    fs.write(js)
+                    fs.write("\n")
+                    fprint(f"-- inputs={se} -- {js}")
+
+    if os.path.exists(agg_stat_file):
+        print(f"-- statistics from {agg_stat_file!r}")
+        import pandas
+
+        df = pandas.read_json(agg_stat_file, lines=True)
+        first = [
+            "timestamp",
+            "model_id",
+            "pretrained",
+            "part",
+            "device",
+            "dtype",
+            "attention",
+            "opset",
+        ]
+        index = [*first[1:], "exporter"]
+        df = df[[*first, *[c for c in df.columns if c not in set(first)]]]
+        df.to_excel(agg_stat_file + ".xlsx")
+
+        values = [
+            "abs",
+            "%>0.1",
+            "%>0.01",
+            "export_duration",
+            "speedup",
+            "latency_torch",
+            "latency_ort_n",
+        ]
+        agg = {
+            **{c: "max" for c in values if c != "speedup"},
+            "speedup": "min",
+        }
+        stat = df[[*index, *values]].groupby(index, dropna=False).agg(agg)
+        stat.to_excel(agg_stat_file + ".agg.xlsx")
+        stat = (
+            df[df.exporter != "custom"][[*index, *values]]
+            .groupby(index, dropna=False)
+            .agg(agg)
+        )
+        stat.to_excel(agg_stat_file + ".agg.onnx-dynamo.xlsx")
+
+
+def zip_model_and_data_into_a_single_file(zip_file: str, model_file: str):
+    """
+    Zips an onnx model and its data into a zingle file.
+
+    :param zip_file: zip file to create
+    :param model_file: onnx file
+    """
+    print()
+    print(f"-- make file {zip_file!r}")
+    cmd = ["zip", "-v", "-1", zip_file]
+    for name in [model_file, f"{model_file}.data"]:
+        print(f"-- add {name!r}")
+        cmd.append(name)
+    print(f"-- cmd: {' '.join(cmd)}")
+    subprocess.run(cmd, check=True)
+    print("-- done.")