onnx-diagnostic 0.8.2__py3-none-any.whl → 0.8.4__py3-none-any.whl

onnx_diagnostic/torch_models/code_sample.py

@@ -236,7 +236,7 @@ def code_sample(
             )
         )
     """
-    model_id, subfolder, same_as_pretrained, use_pretrained = _preprocess_model_id(
+    model_id, subfolder, same_as_pretrained, use_pretrained, submodule = _preprocess_model_id(
         model_id,
         subfolder,
         same_as_pretrained=same_as_pretrained,
@@ -256,6 +256,7 @@ def code_sample(
         model_kwargs=mop,
         subfolder=subfolder,
         add_second_input=False,
+        submodule=submodule,
     )
     if drop_inputs:
         update = {}

onnx_diagnostic/torch_models/hghub/model_inputs.py

@@ -26,17 +26,26 @@ def _code_needing_rewriting(model: Any) -> Any:
 
 
 def _preprocess_model_id(
-    model_id: str, subfolder: Optional[str], same_as_pretrained: bool, use_pretrained: bool
-) -> Tuple[str, Optional[str], bool, bool]:
+    model_id: str,
+    subfolder: Optional[str],
+    same_as_pretrained: bool,
+    use_pretrained: bool,
+    submodule: Optional[str] = None,
+) -> Tuple[str, Optional[str], bool, bool, Optional[str]]:
+    if "::" in model_id:
+        assert (
+            not submodule
+        ), f"submodule={submodule!r} cannot be defined in model_id={model_id!r} as well"
+        model_id, submodule = model_id.split("::", maxsplit=1)
     if subfolder or "//" not in model_id:
-        return model_id, subfolder, same_as_pretrained, use_pretrained
+        return model_id, subfolder, same_as_pretrained, use_pretrained, submodule
     spl = model_id.split("//")
     if spl[-1] == "pretrained":
-        return _preprocess_model_id("//".join(spl[:-1]), "", True, True)
+        return _preprocess_model_id("//".join(spl[:-1]), "", True, True, submodule)
     if spl[-1] in {"transformer", "vae"}:
         # known subfolder
-        return "//".join(spl[:-1]), spl[-1], same_as_pretrained, use_pretrained
-    return model_id, subfolder, same_as_pretrained, use_pretrained
+        return "//".join(spl[:-1]), spl[-1], same_as_pretrained, use_pretrained, submodule
+    return model_id, subfolder, same_as_pretrained, use_pretrained, submodule
 
 
 def get_untrained_model_with_inputs(
@@ -54,6 +63,7 @@ def get_untrained_model_with_inputs(
     subfolder: Optional[str] = None,
     use_only_preinstalled: bool = False,
     config_reduction: Optional[Callable[[Any, str], Dict]] = None,
+    submodule: Optional[str] = None,
 ) -> Dict[str, Any]:
     """
     Gets a non initialized model similar to the original model
@@ -82,6 +92,7 @@ def get_untrained_model_with_inputs(
         <onnx_diagnostic.torch_models.hghub.reduce_model_config>`,
         this function takes a configuration and a task (string)
         as arguments
+    :param submodule: use a submodule instead of the main model
     :return: dictionary with a model, inputs, dynamic shapes, and the configuration,
         some necessary rewriting as well
 
@@ -108,11 +119,12 @@ def get_untrained_model_with_inputs(
         f"model_id={model_id!r}, preinstalled model is only available "
         f"if use_only_preinstalled is False."
     )
-    model_id, subfolder, same_as_pretrained, use_pretrained = _preprocess_model_id(
+    model_id, subfolder, same_as_pretrained, use_pretrained, submodule = _preprocess_model_id(
         model_id,
         subfolder,
         same_as_pretrained=same_as_pretrained,
         use_pretrained=use_pretrained,
+        submodule=submodule,
     )
     if verbose:
         print(
@@ -147,6 +159,8 @@ def get_untrained_model_with_inputs(
         if verbose:
             print(f"[get_untrained_model_with_inputs] architecture={arch!r}")
             print(f"[get_untrained_model_with_inputs] cls={config.__class__.__name__!r}")
+            if submodule:
+                print(f"[get_untrained_model_with_inputs] submodule={submodule!r}")
         if task is None:
             task = task_from_arch(arch, model_id=model_id, subfolder=subfolder)
         if verbose:
@@ -357,6 +371,19 @@ def get_untrained_model_with_inputs(
     if diff_config is not None:
         res["dump_info"] = dict(config_diff=diff_config)
 
+    if submodule:
+        path = submodule.split("::") if "::" in submodule else [submodule]
+        for p in path:
+            assert hasattr(model, p), (
+                f"Unable to find submodule {p!r} in in class {type(model)}, "
+                f"submodule={submodule!r}, possible candidates: "
+                f"{[k for k in dir(model) if isinstance(getattr(model, k), torch.nn.Module)]}"
+            )
+            model = getattr(model, p)
+
+    if verbose:
+        print(f"[get_untrained_model_with_inputs] model class={model.__class__.__name__!r}")
+
     sizes = compute_model_size(model)
     res["model"] = model
     res["configuration"] = config

onnx_diagnostic/torch_models/validate.py

@@ -349,13 +349,15 @@ def _prepare_validation(
     verbose,
     output_names,
     dump_folder,
+    submodule,
 ):
     main_validation_begin = time.perf_counter()
-    model_id, subfolder, same_as_pretrained, use_pretrained = _preprocess_model_id(
+    model_id, subfolder, same_as_pretrained, use_pretrained, submodule = _preprocess_model_id(
         model_id,
         subfolder,
         same_as_pretrained=same_as_pretrained,
         use_pretrained=use_pretrained,
+        submodule=submodule,
     )
     time_preprocess_model_id = time.perf_counter() - main_validation_begin
     patch_kwargs = make_patch_kwargs(patch=patch, rewrite=rewrite)
@@ -364,6 +366,7 @@ def _prepare_validation(
     summary.update(
         dict(
             version_model_id=model_id,
+            version_submodule=submodule,
             version_do_run=str(do_run),
             version_dtype=str(dtype or ""),
             version_device=str(device or ""),
@@ -444,6 +447,7 @@ def _prepare_validation(
         dump_folder,
         folder_name,
         patch_kwargs,
+        submodule,
     )
 
 
@@ -460,6 +464,7 @@ def _get_untrained_model_with_inputs(
     inputs2,
     quiet,
     dump_folder,
+    submodule,
 ):
     iop = input_options or {}
     mop = model_options or {}
@@ -480,6 +485,7 @@ def _get_untrained_model_with_inputs(
                     model_kwargs=mop,
                     subfolder=sub,
                     add_second_input=i2,
+                    submodule=submodule,
                 )
             )
         ),
@@ -671,7 +677,16 @@ def _call_exporter(
     do_run,
     output_names,
     exporter_options,
+    save_ep,
 ):
+    if save_ep and dump_folder:
+        for name in data:
+            if name.startswith("inputs"):
+                if verbose:
+                    print(f"[validate_model] -- dump {name!r}")
+                filename = os.path.join(dump_folder, f"{save_ep}.{name}.pt")
+                torch.save(data[name], filename)
+
     if exporter:
         expop = exporter_options or {}
         if verbose:
@@ -711,6 +726,7 @@ def _call_exporter(
                     dump_folder=dump_folder,
                     output_names=output_names,
                     exporter_options=expop,
+                    save_ep=save_ep,
                 )
         else:
             data["inputs_export"] = data["inputs"]
@@ -831,6 +847,8 @@ def validate_model(
     output_names: Optional[List[str]] = None,
     ort_logs: bool = False,
     quiet_input_sets: Optional[Set[str]] = None,
+    save_ep: Optional[str] = None,
+    submodule: Optional[str] = None,
 ) -> Tuple[Dict[str, Union[int, float, str]], Dict[str, Any]]:
     """
     Validates a model.
@@ -889,6 +907,9 @@ def validate_model(
     :param ort_logs: increases onnxruntime verbosity when creating the session
     :param quiet_input_sets: avoid raising an exception if the inputs belongs to that set
         even if quiet is False
+    :param save_ep: if not empty, this can be used to save the input sets and
+        the exported program
+    :param submodule: to test not the model but a submodule of this model
     :return: two dictionaries, one with some metrics,
         another one with whatever the function produces
 
@@ -952,6 +973,8 @@ def validate_model(
         subfolder=subfolder,
         use_pretrained=use_pretrained,
         same_as_pretrained=same_as_pretrained,
+        save_ep=save_ep,
+        submodule=submodule,
     )
     if dump_folder:
         with open(dump_stats, "w") as f:
@@ -1038,6 +1061,8 @@ def _validate_model_step1(
     subfolder,
     use_pretrained,
     same_as_pretrained,
+    save_ep,
+    submodule,
 ):
     assert not do_same or do_run, (
         f"Discrepancies cannot be measured if the model is not run, "
@@ -1052,6 +1077,7 @@ def _validate_model_step1(
         dump_folder,
         folder_name,
         patch_kwargs,
+        submodule,
     ) = _prepare_validation(
         model_id=model_id,
         subfolder=subfolder,
@@ -1078,6 +1104,7 @@ def _validate_model_step1(
         verbose=verbose,
         output_names=output_names,
         dump_folder=dump_folder,
+        submodule=submodule,
     )
 
     data, iop, mop = _get_untrained_model_with_inputs(
@@ -1093,6 +1120,7 @@ def _validate_model_step1(
         inputs2=inputs2,
         quiet=quiet,
         dump_folder=dump_folder,
+        submodule=submodule,
     )
 
     second_input_keys = [k for k in data if k.startswith("inputs") and k != "inputs"]
@@ -1153,6 +1181,7 @@ def _validate_model_step1(
         do_run=do_run,
         output_names=output_names,
         exporter_options=exporter_options,
+        save_ep=save_ep,
     )
 
     cont, dump_stats = _dump_onnx_model(
@@ -1426,6 +1455,7 @@ def call_exporter(
     dump_folder: Optional[str] = None,
     output_names: Optional[List[str]] = None,
     exporter_options: Optional[Dict[str, Any]] = None,
+    save_ep: Optional[str] = None,
 ) -> Tuple[Dict[str, Union[int, float, str]], Dict[str, Any]]:
     """
     Calls an exporter on a model;
@@ -1440,6 +1470,7 @@ def call_exporter(
     :param dump_folder: to dump additional information
     :param output_names: list of output names to use with the onnx exporter
     :param exporter_options: exporter options
+    :param save_ep: saves the exported program
     :return: two dictionaries, one with some metrics,
         another one with whatever the function produces
     """
@@ -1456,6 +1487,8 @@ def call_exporter(
             optimization=optimization,
             do_run=do_run,
             exporter_options=exporter_options,
+            save_ep=save_ep,
+            dump_folder=dump_folder,
         )
         _restore_torch_export_export(summary)
         return summary, data
@@ -1469,6 +1502,8 @@ def call_exporter(
             optimization=optimization,
             output_names=output_names,
             exporter_options=exporter_options,
+            dump_folder=dump_folder,
+            save_ep=save_ep,
         )
         _restore_torch_export_export(summary)
         return summary, data
@@ -1483,6 +1518,7 @@ def call_exporter(
             dump_folder=dump_folder,
             output_names=output_names,
             exporter_options=exporter_options,
+            save_ep=save_ep,
         )
         _restore_torch_export_export(summary)
         return summary, data
@@ -1516,6 +1552,8 @@ def call_torch_export_export(
     optimization: Optional[str] = None,
     do_run: bool = False,
     exporter_options: Optional[Dict[str, Any]] = None,
+    dump_folder: Optional[str] = None,
+    save_ep: Optional[str] = None,
 ):
     """
     Exports a model with :func:`torch.export.export`.
@@ -1529,6 +1567,8 @@ def call_torch_export_export(
     :param optimization: optimization to do
     :param do_run: runs and compute discrepancies
     :param exporter_options: additional options given to the exporter
+    :param dump_folder: folder where to dump the exported program
+    :param save_ep: to save the exported program
     :return: two dictionaries, one with some metrics,
         another one with whatever the function produces
     """
@@ -1604,6 +1644,12 @@ def call_torch_export_export(
         print(ep)
         print("[call_torch_export_export] -- End of ExportedProgram")
 
+    if dump_folder and save_ep:
+        fname = f"{save_ep}.pt2"
+        if verbose:
+            print(f"[call_torch_export_export] -- save the exported program in {fname!r}")
+        torch.export.save(ep, os.path.join(dump_folder, fname))
+
     if do_run:
         # We check for discrepancies.
         if verbose:
@@ -1880,6 +1926,8 @@ def call_torch_export_onnx(
     optimization: Optional[str] = None,
     output_names: Optional[List[str]] = None,
     exporter_options: Optional[Dict[str, Any]] = None,
+    dump_folder: Optional[str] = None,
+    save_ep: Optional[str] = None,
 ) -> Tuple[Dict[str, Any], Dict[str, Any]]:
     """
     Exports a model into onnx.
@@ -1893,6 +1941,8 @@ def call_torch_export_onnx(
     :param optimization: optimization to do
     :param output_names: output names to use
     :param exporter_options: additional options to give the exporter
+    :param dump_folder: to know where to dump the exported program
+    :param save_ep: to save the exported program
     :return: two dictionaries, one with some metrics,
         another one with whatever the function produces
     """
@@ -1986,6 +2036,12 @@ def call_torch_export_onnx(
         return summary, data
 
     assert epo is not None, "no onnx export was found"
+    if dump_folder and save_ep:
+        fname = f"{save_ep}.pt2"
+        if verbose:
+            print(f"[call_torch_export_export] -- save the exported program in {fname!r}")
+        torch.export.save(epo.exported_program, os.path.join(dump_folder, fname))
+
     if verbose:
         print("[call_torch_export_onnx] done (export)")
     data["onnx_program"] = epo
@@ -2219,6 +2275,7 @@ def call_torch_export_custom(
     dump_folder: Optional[str] = None,
     output_names: Optional[List[str]] = None,
     exporter_options: Optional[Dict[str, Any]] = None,
+    save_ep: Optional[str] = None,
 ) -> Tuple[Dict[str, Any], Dict[str, Any]]:
     """
     Exports a model into onnx.
@@ -2233,6 +2290,7 @@ def call_torch_export_custom(
     :param dump_folder: to store additional information
     :param output_names: list of output names to use
     :param exporter_options: additional exporter options
+    :param save_ep: to save the exported program
     :return: two dictionaries, one with some metrics,
         another one with whatever the function produces
     """
@@ -2345,7 +2403,11 @@ def call_torch_export_custom(
     export_options = ExportOptions(
         strict=strict,
         decomposition_table=decomposition_table,
-        save_ep=(os.path.join(dump_folder, f"{exporter}.ep") if dump_folder else None),
+        save_ep=(
+            (os.path.join(dump_folder, f"{exporter}.ep"), 2**35 if save_ep else 2**18)
+            if dump_folder
+            else None
+        ),
         **exporter_options,
     )
     options = OptimizationOptions(patterns=optimization) if optimization else None

onnx_diagnostic/torch_onnx/runtime_info.py

@@ -4,6 +4,7 @@ import onnx
 import torch
 from ..api import TensorLike
 from ..helpers import string_type
+from ..helpers.onnx_helper import get_hidden_inputs
 
 
 class RuntimeValueKind(enum.IntEnum):
@@ -151,30 +152,6 @@ class RuntimeValue:
         return self.kind == RuntimeValueKind.INITIALIZER
 
 
-def get_hidden_inputs(graph: onnx.GraphProto) -> Set[str]:
-    """
-    Returns the hidden inputs (inputs coming from an upper context)
-    used by a subgraph.
-    """
-    hidden = set()
-    memo = (
-        set(i.name for i in graph.initializer)
-        | set(i.name for i in graph.sparse_initializer)
-        | set(i.name for i in graph.input)
-    )
-    for node in graph.node:
-        for i in node.input:
-            if i not in memo:
-                hidden.add(i)
-        for att in node.attribute:
-            if att.type == onnx.AttributeProto.GRAPH and att.g:
-                hid = get_hidden_inputs(att.g)
-                less = set(h for h in hid if h not in memo)
-                hidden |= less
-        memo |= set(node.output)
-    return hidden
-
-
 def set_is_shape(
     node: onnx.NodeProto, values: Dict[str, RuntimeValue], drop: Optional[Set[str]] = None
 ) -> List[str]: