[ONNX] New export logic leveraging ExportedProgram and ONNX IR (pytorch#132530)

1/n PR to

- Move code from torch-onnx from commit justinchuby/torch-onnx@395495e into torch.onnx and fixes imports.
- Integrate the new export logic with the torch.onnx.export API and include basic set of tests.
- Refactor the API for the change.
- Improve documentation.

Next PRs will be more tests and docs.

Fix pytorch#129277
Pull Request resolved: pytorch#132530
Approved by: https://github.com/titaiwangms, https://github.com/malfet

Author: justinchuby

docs/source/onnx_torchscript.rst

@@ -715,5 +715,5 @@ Classes
     :template: classtemplate.rst
 
     JitScalarType
-    torch.onnx.verification.GraphInfo
-    torch.onnx.verification.VerificationOptions
+    verification.GraphInfo
+    verification.VerificationOptions

mypy.ini

@@ -165,9 +165,6 @@ ignore_missing_imports = True
 [mypy-tensorboard.*]
 ignore_missing_imports = True
 
-[mypy-onnx.*]
-ignore_missing_imports = True
-
 [mypy-matplotlib.*]
 ignore_missing_imports = True
 
@@ -301,5 +298,14 @@ ignore_missing_imports = True
 # Third party dependencies that are optional.
 #
 
+[mypy-onnx.*]
+ignore_missing_imports = True
+
+[mypy-onnxruntime.*]
+ignore_missing_imports = True
+
+[mypy-onnxscript.*]
+ignore_missing_imports = True
+
 [mypy-redis]
 ignore_missing_imports = True

test/onnx/dynamo/test_exporter_api.py

@@ -163,222 +163,5 @@ def test_raise_from_diagnostic_warning_when_diagnostic_option_warning_as_error_i
             )
 
 
-class TestONNXExportWithDynamo(common_utils.TestCase):
-    def test_args_normalization_with_no_kwargs(self):
-        exported_program = torch.export.export(
-            SampleModelTwoInputs(),
-            (
-                torch.randn(1, 1, 2),
-                torch.randn(1, 1, 2),
-            ),
-        )
-        onnx_program_from_new_exporter = torch.onnx.dynamo_export(
-            exported_program, torch.randn(1, 1, 2), torch.randn(1, 1, 2)
-        )
-        onnx_program_from_old_exporter = torch.onnx.export(
-            SampleModelTwoInputs(),
-            (torch.randn(1, 1, 2), torch.randn(1, 1, 2)),
-            dynamo=True,
-        )
-        self.assertEqual(
-            onnx_program_from_new_exporter.model_proto,
-            onnx_program_from_old_exporter.model_proto,
-        )
-
-    def test_args_is_tensor_not_tuple(self):
-        exported_program = torch.export.export(SampleModel(), (torch.randn(1, 1, 2),))
-        onnx_program_from_new_exporter = torch.onnx.dynamo_export(
-            exported_program, torch.randn(1, 1, 2)
-        )
-        onnx_program_from_old_exporter = torch.onnx.export(
-            SampleModel(), torch.randn(1, 1, 2), dynamo=True
-        )
-        self.assertEqual(
-            onnx_program_from_new_exporter.model_proto,
-            onnx_program_from_old_exporter.model_proto,
-        )
-
-    def test_args_normalization_with_kwargs(self):
-        exported_program = torch.export.export(
-            SampleModelTwoInputs(), (torch.randn(1, 1, 2),), {"b": torch.randn(1, 1, 2)}
-        )
-        onnx_program_from_new_exporter = torch.onnx.dynamo_export(
-            exported_program, torch.randn(1, 1, 2), b=torch.randn(1, 1, 2)
-        )
-        onnx_program_from_old_exporter = torch.onnx.export(
-            SampleModelTwoInputs(),
-            (torch.randn(1, 1, 2), {"b": torch.randn(1, 1, 2)}),
-            dynamo=True,
-        )
-        self.assertEqual(
-            onnx_program_from_new_exporter.model_proto,
-            onnx_program_from_old_exporter.model_proto,
-        )
-
-    def test_args_normalization_with_empty_dict_at_the_tail(self):
-        exported_program = torch.export.export(
-            SampleModelTwoInputs(), (torch.randn(1, 1, 2),), {"b": torch.randn(1, 1, 2)}
-        )
-        onnx_program_from_new_exporter = torch.onnx.dynamo_export(
-            exported_program, torch.randn(1, 1, 2), b=torch.randn(1, 1, 2)
-        )
-        onnx_program_from_old_exporter = torch.onnx.export(
-            SampleModelTwoInputs(),
-            (torch.randn(1, 1, 2), {"b": torch.randn(1, 1, 2)}),
-            dynamo=True,
-        )
-        self.assertEqual(
-            onnx_program_from_new_exporter.model_proto,
-            onnx_program_from_old_exporter.model_proto,
-        )
-
-    def test_dynamic_axes_enable_dynamic_shapes_with_fully_specified_axes(self):
-        exported_program = torch.export.export(
-            SampleModelForDynamicShapes(),
-            (
-                torch.randn(2, 2, 3),
-                torch.randn(2, 2, 3),
-            ),
-            dynamic_shapes={
-                "x": {
-                    0: torch.export.Dim("customx_dim_0"),
-                    1: torch.export.Dim("customx_dim_1"),
-                    2: torch.export.Dim("customx_dim_2"),
-                },
-                "b": {
-                    0: torch.export.Dim("customb_dim_0"),
-                    1: torch.export.Dim("customb_dim_1"),
-                    2: torch.export.Dim("customb_dim_2"),
-                },
-            },
-        )
-        onnx_program_from_new_exporter = torch.onnx.dynamo_export(
-            exported_program,
-            torch.randn(2, 2, 3),
-            b=torch.randn(2, 2, 3),
-        )
-        onnx_program_from_old_exporter = torch.onnx.export(
-            SampleModelForDynamicShapes(),
-            (torch.randn(2, 2, 3), {"b": torch.randn(2, 2, 3)}),
-            dynamic_axes={
-                "x": {0: "customx_dim_0", 1: "customx_dim_1", 2: "customx_dim_2"},
-                "b": {0: "customb_dim_0", 1: "customb_dim_1", 2: "customb_dim_2"},
-            },
-            dynamo=True,
-        )
-        self.assertEqual(
-            onnx_program_from_new_exporter.model_proto,
-            onnx_program_from_old_exporter.model_proto,
-        )
-
-    def test_dynamic_axes_enable_dynamic_shapes_with_default_axe_names(self):
-        exported_program = torch.export.export(
-            SampleModelForDynamicShapes(),
-            (
-                torch.randn(2, 2, 3),
-                torch.randn(2, 2, 3),
-            ),
-            dynamic_shapes={
-                "x": {
-                    0: torch.export.Dim("customx_dim_0"),
-                    1: torch.export.Dim("customx_dim_1"),
-                    2: torch.export.Dim("customx_dim_2"),
-                },
-                "b": {
-                    0: torch.export.Dim("customb_dim_0"),
-                    1: torch.export.Dim("customb_dim_1"),
-                    2: torch.export.Dim("customb_dim_2"),
-                },
-            },
-        )
-        onnx_program_from_new_exporter = torch.onnx.dynamo_export(
-            exported_program,
-            torch.randn(2, 2, 3),
-            b=torch.randn(2, 2, 3),
-        )
-        onnx_program_from_old_exporter = torch.onnx.export(
-            SampleModelForDynamicShapes(),
-            (torch.randn(2, 2, 3), {"b": torch.randn(2, 2, 3)}),
-            dynamic_axes={
-                "x": [0, 1, 2],
-                "b": [0, 1, 2],
-            },
-            dynamo=True,
-        )
-        self.assertEqual(
-            onnx_program_from_new_exporter.model_proto,
-            onnx_program_from_old_exporter.model_proto,
-        )
-
-    def test_dynamic_axes_supports_partial_dynamic_shapes(self):
-        exported_program = torch.export.export(
-            SampleModelForDynamicShapes(),
-            (
-                torch.randn(2, 2, 3),
-                torch.randn(2, 2, 3),
-            ),
-            dynamic_shapes={
-                "x": None,
-                "b": {
-                    0: torch.export.Dim("customb_dim_0"),
-                    1: torch.export.Dim("customb_dim_1"),
-                    2: torch.export.Dim("customb_dim_2"),
-                },
-            },
-        )
-        onnx_program_from_new_exporter = torch.onnx.dynamo_export(
-            exported_program,
-            torch.randn(2, 2, 3),
-            b=torch.randn(2, 2, 3),
-        )
-        onnx_program_from_old_exporter = torch.onnx.export(
-            SampleModelForDynamicShapes(),
-            (torch.randn(2, 2, 3), {"b": torch.randn(2, 2, 3)}),
-            dynamic_axes={
-                "b": [0, 1, 2],
-            },
-            dynamo=True,
-        )
-        self.assertEqual(
-            onnx_program_from_new_exporter.model_proto,
-            onnx_program_from_old_exporter.model_proto,
-        )
-
-    def test_dynamic_shapes_hit_constraints_in_dynamo(self):
-        # SampleModelTwoInputs has constraints becuse of add of two inputs,
-        # so the two input shapes are related.
-        with self.assertRaisesRegex(
-            torch._dynamo.exc.UserError,
-            "Constraints violated",
-        ):
-            _ = torch.onnx.export(
-                SampleModelTwoInputs(),
-                (torch.randn(2, 2, 3), torch.randn(2, 2, 3)),
-                dynamic_axes={
-                    "x": {0: "x_dim_0", 1: "x_dim_1", 2: "x_dim_2"},
-                    "b": {0: "b_dim_0", 1: "b_dim_1", 2: "b_dim_2"},
-                },
-                dynamo=True,
-            )
-
-    def test_saved_f_exists_after_export(self):
-        with common_utils.TemporaryFileName(suffix=".onnx") as path:
-            _ = torch.onnx.export(
-                SampleModel(), torch.randn(1, 1, 2), path, dynamo=True
-            )
-            self.assertTrue(os.path.exists(path))
-
-    def test_raises_error_when_input_is_script_module(self):
-        class ScriptModule(torch.jit.ScriptModule):
-            def forward(self, x):
-                return x
-
-        with self.assertRaisesRegex(
-            TypeError,
-            "Dynamo export does not support ScriptModule or ScriptFunction.",
-        ):
-            _ = torch.onnx.export(ScriptModule(), torch.randn(1, 1, 2), dynamo=True)
-
-
 if __name__ == "__main__":
     common_utils.run_tests()

test/onnx/exporter/README.md

@@ -0,0 +1 @@
+Directory for all ExportedProgram exporter logic.

test/onnx/exporter/test_api.py

@@ -0,0 +1,120 @@
+# Owner(s): ["module: onnx"]
+"""Simple API tests for the ONNX exporter."""
+
+from __future__ import annotations
+
+import os
+
+import torch
+from torch.onnx._internal import exporter
+from torch.testing._internal import common_utils
+
+
+class SampleModel(torch.nn.Module):
+    def forward(self, x):
+        y = x + 1
+        z = y.relu()
+        return (y, z)
+
+
+class SampleModelTwoInputs(torch.nn.Module):
+    def forward(self, x, b):
+        y = x + b
+        z = y.relu()
+        return (y, z)
+
+
+class SampleModelForDynamicShapes(torch.nn.Module):
+    def forward(self, x, b):
+        return x.relu(), b.sigmoid()
+
+
+class TestExportAPIDynamo(common_utils.TestCase):
+    """Tests for the ONNX exporter API when dynamo=True."""
+
+    def test_args_normalization_with_no_kwargs(self):
+        onnx_program = torch.onnx.export(
+            SampleModelTwoInputs(),
+            (torch.randn(1, 1, 2), torch.randn(1, 1, 2)),
+            dynamo=True,
+        )
+        assert onnx_program
+        exporter.verify_onnx_program(onnx_program)
+
+    def test_args_normalization_with_kwargs(self):
+        onnx_program = torch.onnx.export(
+            SampleModelTwoInputs(),
+            (torch.randn(1, 1, 2), {"b": torch.randn(1, 1, 2)}),
+            dynamo=True,
+        )
+        assert onnx_program
+        exporter.verify_onnx_program(onnx_program)
+
+    def test_args_normalization_with_empty_dict_at_the_tail(self):
+        onnx_program = torch.onnx.export(
+            SampleModelTwoInputs(),
+            (torch.randn(1, 1, 2), {"b": torch.randn(1, 1, 2)}),
+            dynamo=True,
+        )
+        assert onnx_program
+        exporter.verify_onnx_program(onnx_program)
+
+    def test_dynamic_axes_enable_dynamic_shapes_with_fully_specified_axes(self):
+        onnx_program = torch.onnx.export(
+            SampleModelForDynamicShapes(),
+            (torch.randn(2, 2, 3), {"b": torch.randn(2, 2, 3)}),
+            dynamic_axes={
+                "x": {0: "customx_dim_0", 1: "customx_dim_1", 2: "customx_dim_2"},
+                "b": {0: "customb_dim_0", 1: "customb_dim_1", 2: "customb_dim_2"},
+            },
+            dynamo=True,
+        )
+        assert onnx_program
+        exporter.verify_onnx_program(onnx_program)
+
+    def test_dynamic_axes_enable_dynamic_shapes_with_default_axe_names(self):
+        onnx_program = torch.onnx.export(
+            SampleModelForDynamicShapes(),
+            (torch.randn(2, 2, 3), {"b": torch.randn(2, 2, 3)}),
+            dynamic_axes={
+                "x": [0, 1, 2],
+                "b": [0, 1, 2],
+            },
+            dynamo=True,
+        )
+        assert onnx_program
+        exporter.verify_onnx_program(onnx_program)
+
+    def test_dynamic_axes_supports_partial_dynamic_shapes(self):
+        onnx_program = torch.onnx.export(
+            SampleModelForDynamicShapes(),
+            (torch.randn(2, 2, 3), {"b": torch.randn(2, 2, 3)}),
+            dynamic_axes={
+                "b": [0, 1, 2],
+            },
+            dynamo=True,
+        )
+        assert onnx_program
+        exporter.verify_onnx_program(onnx_program)
+
+    def test_saved_f_exists_after_export(self):
+        with common_utils.TemporaryFileName(suffix=".onnx") as path:
+            _ = torch.onnx.export(
+                SampleModel(), (torch.randn(1, 1, 2),), path, dynamo=True
+            )
+            self.assertTrue(os.path.exists(path))
+
+    def test_export_supports_script_module(self):
+        class ScriptModule(torch.nn.Module):
+            def forward(self, x):
+                return x
+
+        onnx_program = torch.onnx.export(
+            torch.jit.script(ScriptModule()), (torch.randn(1, 1, 2),), dynamo=True
+        )
+        assert onnx_program
+        exporter.verify_onnx_program(onnx_program)
+
+
+if __name__ == "__main__":
+    common_utils.run_tests()