replay onnx with torch

Author: xadupre

_unittests/ut_torch_onnx/test_sbs.py

@@ -379,7 +379,7 @@ def forward(self, x):
                 use_tensor=True,
             ),
         )
-        df = pandas.DataFrame(list(results))
+        df = pandas.DataFrame(list(results)).dropna(axis=1, how="all")
         df.to_excel(self.get_dump_file("test_sbs_model_with_weights_custom.xlsx"))
         self.assertEqual(
             [
@@ -391,7 +391,6 @@ def forward(self, x):
                 "err_abs",
                 "err_dev",
                 "err_h01",
-                "err_nan",
                 "err_rel",
                 "onnx_id_node",
                 "onnx_id_output",
@@ -445,7 +444,7 @@ def forward(self, x):
                 use_tensor=True,
             ),
         )
-        df = pandas.DataFrame(list(results))
+        df = pandas.DataFrame(list(results)).dropna(axis=1, how="all")
         df.to_excel(self.get_dump_file("test_sbs_model_with_weights_dynamo.xlsx"))
         self.assertEqual(
             [
@@ -457,7 +456,6 @@ def forward(self, x):
                 "err_abs",
                 "err_dev",
                 "err_h01",
-                "err_nan",
                 "err_rel",
                 "onnx_id_node",
                 "onnx_id_output",
@@ -542,7 +540,7 @@ def forward(self, x):
                 reset_names=["linear"],
             ),
         )
-        df = pandas.DataFrame(list(results))
+        df = pandas.DataFrame(list(results)).dropna(axis=1, how="all")
         df.to_excel(self.get_dump_file("test_sbs_model_with_weights_custom_reset.xlsx"))
         onnx_op_type = df["onnx_op_type"].tolist()
         self.assertEqual(onnx_op_type.count("reset"), 1)
@@ -593,10 +591,80 @@ def forward(self, x):
                 ),
             ),
         )
-        df = pandas.DataFrame(list(results))
+        df = pandas.DataFrame(list(results)).dropna(axis=1, how="all")
         df.to_excel(self.get_dump_file("test_sbs_replay.xlsx"))
-        print(df)
-        # self.clean_dump()
+        self.assertEqual(df.shape, (8, 15))
+        self.clean_dump()
+
+    @hide_stdout()
+    @ignore_warnings((DeprecationWarning, FutureWarning, UserWarning))
+    def test_sbs_run_onnx_with_torch_inputs(self):
+        torch = self.torch
+
+        class Model(self.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=self.torch.randn((5, 10)))
+        ds = dict(x={0: "batch"})
+        Model()(**inputs)
+        ep = self.torch.export.export(
+            Model(), (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds)
+        )
+        filename = self.get_dump_file("test_sbs_run_onnx_with_torch_inputs.onnx")
+        to_onnx(ep, exporter="custom", filename=filename)
+        onx = onnx.load(filename)
+        results = list(
+            run_aligned(
+                ep,
+                onx,
+                kwargs=inputs,
+                run_cls=OnnxruntimeEvaluator,
+                verbose=11,
+                use_tensor=True,
+                run_onnx_with_torch_inputs=True,
+            ),
+        )
+        df = pandas.DataFrame(list(results)).dropna(axis=1, how="all")
+        df.to_excel(self.get_dump_file("test_sbs_run_onnx_with_torch_inputs.xlsx"))
+        self.assertEqual(
+            [
+                "ep_id_node",
+                "ep_name",
+                "ep_shape_type",
+                "ep_target",
+                "ep_time_run",
+                "err_abs",
+                "err_abs2",
+                "err_dev",
+                "err_dev2",
+                "err_h01",
+                "err_h012",
+                "err_rel",
+                "err_rel2",
+                "onnx_id_node",
+                "onnx_id_output",
+                "onnx_name",
+                "onnx_op_type",
+                "onnx_shape_type",
+                "onnx_time_run",
+            ],
+            sorted(df.columns),
+        )
+        self.assertEqual(len(results), 8)
+        self.assertEqual([0, 0, 0, 0, None, 0, 0, 0], [r.err_dev for r in results])
+        self.assertEqual(
+            [-1, -1, -1, -1, -1, 0, 1, 2], df["onnx_id_node"].fillna(-10).tolist()
+        )
+        self.clean_dump()
 
 
 if __name__ == "__main__":

onnx_diagnostic/_command_lines_parser.py

@@ -1217,6 +1217,19 @@ def get_parser_sbs() -> ArgumentParser:
         default=False,
         help="First runs the whole model.",
     )
+    parser.add_argument(
+        "-2",
+        "--second-run",
+        action=BooleanOptionalAction,
+        default=False,
+        help=textwrap.dedent(
+            """
+            Tries to run all onnx nodes with torch results produced by the exported
+            program. It then measures the discrepancies again. It can be used
+            to identify kernel introduces discrepancies from other just propagating them.
+            """
+        ),
+    )
     parser.add_argument(
         "--reset",
         required=False,
@@ -1365,6 +1378,7 @@ def _size(name):
         reset_names=args.reset.split(","),
         exc=False,
         replay_configuration=replay_configuration,
+        run_onnx_with_torch_inputs=args.second_run,
     ):
         data.append(obs)
         if (
@@ -1377,8 +1391,10 @@ def _size(name):
             )
             df.to_excel(args.output)
     print(f"-- final saves into {args.output!r}")
-    df = pandas.DataFrame(data).apply(
-        lambda col: col.fillna("") if col.dtype == "object" else col
+    df = (
+        pandas.DataFrame(data)
+        .apply(lambda col: col.fillna("") if col.dtype == "object" else col)
+        .dropna(axis=1, how="all")
     )
     df.to_excel(args.output, index=False)
     print("-- done")