torch-onnx

Author: xadupre

CHANGELOGS.rst

@@ -4,7 +4,8 @@ Change Logs
 0.6.1
 +++++
 
-* :pr:`115`, :pr:`116`, :pr:`117`, :pr:`118`: first steps for TorchOnnxEvaluator
+* :pr:`115`, :pr:`116`, :pr:`117`, :pr:`118`, :pr:`119`:
+  first steps for TorchOnnxEvaluator
 * :pr:`114`: extends the list of known rewritings
 * :pr:`113`: fixes a couple of issues with ModelBuilder
 

_unittests/ut_torch_models/test_test_helpers.py

@@ -164,6 +164,34 @@ def test_validate_model_custom(self):
             onnx_filename, output_path, num_attention_heads=2, hidden_size=192, verbose=10
         )
 
+    @requires_torch("2.7")
+    @hide_stdout()
+    @ignore_warnings(FutureWarning)
+    @requires_experimental()
+    def test_validate_model_custom_torch(self):
+        mid = "arnir0/Tiny-LLM"
+        summary, data = validate_model(
+            mid,
+            do_run=True,
+            verbose=10,
+            exporter="custom",
+            dump_folder="dump_test_validate_model_custom_torch",
+            patch=True,
+            stop_if_static=2 if pv.Version(torch.__version__) > pv.Version("2.6.1") else 0,
+            optimization="default",
+            quiet=False,
+            runtime="torch",
+        )
+        self.assertIsInstance(summary, dict)
+        self.assertIsInstance(data, dict)
+        self.assertIn("disc_onnx_ort_run_abs", summary)
+        self.assertLess(summary["disc_onnx_ort_run_abs"], 1e-4)
+        onnx_filename = data["onnx_filename"]
+        output_path = f"{onnx_filename}.ortopt.onnx"
+        run_ort_fusion(
+            onnx_filename, output_path, num_attention_heads=2, hidden_size=192, verbose=10
+        )
+
     def test_filter_inputs(self):
         inputs, ds = {"a": 1, "b": 2}, {"a": 20, "b": 30}
         ni, nd = filter_inputs(inputs, dynamic_shapes=ds, drop_names=["a"])

onnx_diagnostic/_command_lines_parser.py

@@ -352,6 +352,12 @@ def get_parser_validate() -> ArgumentParser:
         action=BooleanOptionalAction,
         help="validate the trained model (requires downloading)",
     )
+    parser.add_argument(
+        "--runtime",
+        choices=["onnxruntime", "torch", "ref"],
+        default="onnxruntime",
+        help="onnx runtime to use, ",
+    )
     parser.add_argument(
         "-o",
         "--dump-folder",
@@ -453,6 +459,7 @@ def _cmd_validate(argv: List[Any]):
             model_options=args.mop,
             subfolder=args.subfolder,
             opset=args.opset,
+            runtime=args.runtime,
         )
         print("")
         print("-- summary --")

onnx_diagnostic/reference/torch_evaluator.py

@@ -60,6 +60,8 @@ def __init__(
         else:
             self.default_device = self.CPU
 
+        if isinstance(proto, str):
+            proto = onnx.load(proto)
         if isinstance(proto, onnx.ModelProto):
             assert opsets is None, "proto is a model, opsets must be None in that case"
             assert not proto.graph.sparse_initializer, "sparse_initializer not support yet"

onnx_diagnostic/torch_models/test_helper.py

@@ -17,6 +17,7 @@
 from ..tasks import random_input_kwargs
 from ..torch_export_patches import torch_export_patches
 from ..torch_export_patches.patch_inputs import use_dyn_not_str
+from ..reference import TorchOnnxEvaluator
 from .hghub import get_untrained_model_with_inputs
 
 
@@ -244,6 +245,7 @@ def validate_model(
     model_options: Optional[Dict[str, Any]] = None,
     subfolder: Optional[str] = None,
     opset: Optional[int] = None,
+    runtime: str = "onnxruntime",
 ) -> Tuple[Dict[str, Union[int, float, str]], Dict[str, Any]]:
     """
     Validates a model.
@@ -280,6 +282,8 @@ def validate_model(
         ``num_hidden_layers`` or ``attn_implementation``
     :param subfolder: version or subfolders to uses when retrieving a model id
     :param opset: onnx opset to use for the conversion
+    :param runtime: onnx runtime to use to check about discrepancies,
+        only if `do_run` is true
     :return: two dictionaries, one with some metrics,
         another one with whatever the function produces
 
@@ -308,6 +312,7 @@ def validate_model(
             version_ortfusiontype=ortfusiontype or "",
             version_stop_if_static=str(stop_if_static),
             version_exporter=exporter or "",
+            version_runtime=runtime,
         )
     )
     if opset:
@@ -633,7 +638,9 @@ def validate_model(
         return summary, data
 
     if do_run:
-        summary_valid, data = validate_onnx_model(data=data, quiet=quiet, verbose=verbose)
+        summary_valid, data = validate_onnx_model(
+            data=data, quiet=quiet, verbose=verbose, runtime=runtime
+        )
         summary.update(summary_valid)
 
     if ortfusiontype and "onnx_filename" in data:
@@ -686,7 +693,7 @@ def validate_model(
 
             if do_run:
                 summary_valid, data = validate_onnx_model(
-                    data=data, quiet=quiet, verbose=verbose, flavour=flavour
+                    data=data, quiet=quiet, verbose=verbose, flavour=flavour, runtime=runtime
                 )
                 summary.update(summary_valid)
 
@@ -898,6 +905,7 @@ def validate_onnx_model(
     quiet: bool = False,
     verbose: int = 0,
     flavour: Optional[str] = None,
+    runtime: str = "onnxruntime",
 ) -> Tuple[Dict[str, Any], Dict[str, Any]]:
     """
     Verifies that an onnx model produces the same
@@ -910,6 +918,7 @@ def validate_onnx_model(
     :param quiet: catch exception or not
     :param verbose: verbosity
     :param flavour: use a different version of the inputs
+    :param runtime: onnx runtime to use, onnxruntime or torch
     :return: two dictionaries, one with some metrics,
         another one with whatever the function produces
     """
@@ -951,16 +960,23 @@ def _mk(key):
             f"{providers}..., flavour={flavour!r}"
         )
 
+    assert runtime == "torch", f"runtime={runtime!r}"
+    cls_runtime = (
+        (
+            lambda model, providers: onnxruntime.InferenceSession(
+                (model.SerializeToString() if isinstance(model, onnx.ModelProto) else model),
+                providers=providers,
+            )
+        )
+        if runtime == "onnxruntime"
+        else (lambda model, providers: TorchOnnxEvaluator(model, providers=providers))
+    )
     sess = _quiet_or_not_quiet(
         quiet,
         _mk("time_onnx_ort_create"),
         summary,
         data,
-        (
-            lambda source=source, providers=providers: onnxruntime.InferenceSession(
-                source, providers=providers
-            )
-        ),
+        (lambda source=source, providers=providers: cls_runtime(source, providers)),
     )
     if f"ERR_{_mk('time_onnx_ort_create')}" in summary:
         return summary, data