[ONNX] Introduce Input/Ouptut adapter; Switch to 'DynamoExporter' #98421

BowenBao · 2023-04-05T16:21:28Z

Stack from ghstack (oldest at bottom):

Summary

Introduce input/output adapter. Due to design differences, input/output format
between PyTorch model and exported ONNX model are often not the same. E.g., None
inputs are allowed for PyTorch model, but are not supported by ONNX. Nested constructs
of tensors are allowed for PyTorch model, but only flattened tensors are supported by ONNX,
etc. The new input/output adapter is exported with the model. Providing an interface to
automatically convert and validate inputs/outputs format.
As suggested by [Dynamo] Enable dynamo.export for huggingface models w/ ModelOutput #98251,
provide extension for unwrapping user defined python classes for dynamo.export based
exporter. Unblock huggingface models.
Re-wire tests to run through DynamoExporter w/ dynamo_export api. Kept
DynamoOptimizeExporter in the tests for now for coverage of this change.

[ghstack-poisoned]

pytorch-bot · 2023-04-05T16:21:33Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/98421

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

⏳ No Failures, 2 Pending

As of commit 73a01d0:
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

ghstack-source-id: 39cc6e28980ebd06e249ec3c37e55330434dd2b6 Pull Request resolved: #98421

torch/onnx/_internal/exporter.py

torch/onnx/_internal/fx/dynamo_exporter.py

…matter; Switch to 'DynamoExporter'" Summary * Introduce input/output formatter. Due to design differences, input/output format between PyTorch model and exported ONNX model are often not the same. E.g., `None` inputs are allowed for PyTorch model, but are not supported by ONNX. Nested constructs of tensors are allowed for PyTorch model, but only flattened tensors are supported by ONNX, etc. The new input/output formatter is exported with the model. Providing an interface to automatically convert and validate inputs/outputs format. * As suggested by #98251, provide extension for unwrapping user defined python classes for `dynamo.export` based exporter. Unblock huggingface models. * Re-wire tests to run through `DynamoExporter` w/ `dynamo_export` api. Kept `DynamoOptimizeExporter` in the tests for now for coverage of this change. [ghstack-poisoned]

ghstack-source-id: 15b53c15756e096c844f01764ebae0a55ca57996 Pull Request resolved: #98421

torch/onnx/_internal/exporter.py

…t formatter; Switch to 'DynamoExporter'" Summary * Introduce input/output formatter. Due to design differences, input/output format between PyTorch model and exported ONNX model are often not the same. E.g., `None` inputs are allowed for PyTorch model, but are not supported by ONNX. Nested constructs of tensors are allowed for PyTorch model, but only flattened tensors are supported by ONNX, etc. The new input/output formatter is exported with the model. Providing an interface to automatically convert and validate inputs/outputs format. * As suggested by #98251, provide extension for unwrapping user defined python classes for `dynamo.export` based exporter. Unblock huggingface models. * Re-wire tests to run through `DynamoExporter` w/ `dynamo_export` api. Kept `DynamoOptimizeExporter` in the tests for now for coverage of this change. [ghstack-poisoned]

ghstack-source-id: 35aa14ab11fd80f46557880e3f30fbcc686d7bf4 Pull Request resolved: #98421

…uptut formatter; Switch to 'DynamoExporter'" Summary * Introduce input/output formatter. Due to design differences, input/output format between PyTorch model and exported ONNX model are often not the same. E.g., `None` inputs are allowed for PyTorch model, but are not supported by ONNX. Nested constructs of tensors are allowed for PyTorch model, but only flattened tensors are supported by ONNX, etc. The new input/output formatter is exported with the model. Providing an interface to automatically convert and validate inputs/outputs format. * As suggested by #98251, provide extension for unwrapping user defined python classes for `dynamo.export` based exporter. Unblock huggingface models. * Re-wire tests to run through `DynamoExporter` w/ `dynamo_export` api. Kept `DynamoOptimizeExporter` in the tests for now for coverage of this change. [ghstack-poisoned]

ghstack-source-id: 417f3a993252b55823ff6597dd6817a0c663b5f0 Pull Request resolved: #98421

… Switch to 'DynamoExporter'" Summary * Introduce input/output formatter. Due to design differences, input/output format between PyTorch model and exported ONNX model are often not the same. E.g., `None` inputs are allowed for PyTorch model, but are not supported by ONNX. Nested constructs of tensors are allowed for PyTorch model, but only flattened tensors are supported by ONNX, etc. The new input/output formatter is exported with the model. Providing an interface to automatically convert and validate inputs/outputs format. * As suggested by #98251, provide extension for unwrapping user defined python classes for `dynamo.export` based exporter. Unblock huggingface models. * Re-wire tests to run through `DynamoExporter` w/ `dynamo_export` api. Kept `DynamoOptimizeExporter` in the tests for now for coverage of this change. [ghstack-poisoned]

ghstack-source-id: f7d6e5b42843e96f93a196b31a34623d45aba8f1 Pull Request resolved: #98421

…namoExporter'" Summary * Introduce input/output formatter. Due to design differences, input/output format between PyTorch model and exported ONNX model are often not the same. E.g., `None` inputs are allowed for PyTorch model, but are not supported by ONNX. Nested constructs of tensors are allowed for PyTorch model, but only flattened tensors are supported by ONNX, etc. The new input/output formatter is exported with the model. Providing an interface to automatically convert and validate inputs/outputs format. * As suggested by #98251, provide extension for unwrapping user defined python classes for `dynamo.export` based exporter. Unblock huggingface models. * Re-wire tests to run through `DynamoExporter` w/ `dynamo_export` api. Kept `DynamoOptimizeExporter` in the tests for now for coverage of this change. [ghstack-poisoned]

ghstack-source-id: 24cf566f4a1dfff0e3828f63fc44168ec983b0ae Pull Request resolved: #98421

…Exporter'" Summary * Introduce input/output formatter. Due to design differences, input/output format between PyTorch model and exported ONNX model are often not the same. E.g., `None` inputs are allowed for PyTorch model, but are not supported by ONNX. Nested constructs of tensors are allowed for PyTorch model, but only flattened tensors are supported by ONNX, etc. The new input/output formatter is exported with the model. Providing an interface to automatically convert and validate inputs/outputs format. * As suggested by #98251, provide extension for unwrapping user defined python classes for `dynamo.export` based exporter. Unblock huggingface models. * Re-wire tests to run through `DynamoExporter` w/ `dynamo_export` api. Kept `DynamoOptimizeExporter` in the tests for now for coverage of this change. [ghstack-poisoned]

ghstack-source-id: cfd74b3d586468ec1c026c8c3dc62fb50ff6cae5 Pull Request resolved: #98421

…Exporter'" Summary * Introduce input/output formatter. Due to design differences, input/output format between PyTorch model and exported ONNX model are often not the same. E.g., `None` inputs are allowed for PyTorch model, but are not supported by ONNX. Nested constructs of tensors are allowed for PyTorch model, but only flattened tensors are supported by ONNX, etc. The new input/output formatter is exported with the model. Providing an interface to automatically convert and validate inputs/outputs format. * As suggested by #98251, provide extension for unwrapping user defined python classes for `dynamo.export` based exporter. Unblock huggingface models. * Re-wire tests to run through `DynamoExporter` w/ `dynamo_export` api. Kept `DynamoOptimizeExporter` in the tests for now for coverage of this change. [ghstack-poisoned]

…er; Switch to 'DynamoExporter'" Summary * Introduce input/output formatter. Due to design differences, input/output format between PyTorch model and exported ONNX model are often not the same. E.g., `None` inputs are allowed for PyTorch model, but are not supported by ONNX. Nested constructs of tensors are allowed for PyTorch model, but only flattened tensors are supported by ONNX, etc. The new input/output formatter is exported with the model. Providing an interface to automatically convert and validate inputs/outputs format. * As suggested by #98251, provide extension for unwrapping user defined python classes for `dynamo.export` based exporter. Unblock huggingface models. * Re-wire tests to run through `DynamoExporter` w/ `dynamo_export` api. Kept `DynamoOptimizeExporter` in the tests for now for coverage of this change. [ghstack-poisoned]

ghstack-source-id: 8ff147e1094f5943ed04dd462be88c42d892df08 Pull Request resolved: #98421

test/onnx/test_fx_to_onnx_with_onnxruntime.py

…duce Input/Ouptut formatter; Switch to 'DynamoExporter'" Summary * Introduce input/output formatter. Due to design differences, input/output format between PyTorch model and exported ONNX model are often not the same. E.g., `None` inputs are allowed for PyTorch model, but are not supported by ONNX. Nested constructs of tensors are allowed for PyTorch model, but only flattened tensors are supported by ONNX, etc. The new input/output formatter is exported with the model. Providing an interface to automatically convert and validate inputs/outputs format. * As suggested by #98251, provide extension for unwrapping user defined python classes for `dynamo.export` based exporter. Unblock huggingface models. * Re-wire tests to run through `DynamoExporter` w/ `dynamo_export` api. Kept `DynamoOptimizeExporter` in the tests for now for coverage of this change. [ghstack-poisoned]

ghstack-source-id: c89f1cda4c217bb1cd593c13993eedf0c1e72b6a Pull Request resolved: #98421

thiagocrepaldi · 2023-04-14T18:47:49Z

test/onnx/test_fx_to_onnx_with_onnxruntime.py

+            DynamoExporter: (
+                "beartype.roar.BeartypeCallHintReturnViolation: @beartyped "
+                "torch.onnx._internal.exporter.ExportOutput.adapt_torch_inputs_to_onnx() "
+                "return (tensor([[[ 1.5410, -0.2934]]]), 8.0) violates type hint "


Could we change type Hint to Union[torch.Tensor, int, float, bool]and have these built-in converted to constant tensors during onnx export?

We can, it can happen in follow ups though. I'm unsure yet if we'd need adjustments on fx_to_onnx conversion pass, so don't want to extend the scope of this PR.

thiagocrepaldi · 2023-04-14T18:48:14Z

test/onnx/test_fx_to_onnx_with_onnxruntime.py

+                "instance of <protocol 'torch.Tensor'>."
+            ),
+            DynamoOptimizeExporter: (
+                "RuntimeError: The two modules have different number of arguments. "


Is there a reason the input/output adapter cannot handle this?

Same as the other reply.

test/onnx/test_fx_to_onnx_with_onnxruntime.py

thiagocrepaldi · 2023-04-14T19:02:38Z

torch/onnx/_internal/fx/dynamo_exporter.py

+        return self.export_fx_to_onnx(compiler.captured_graph, model_args)
+
+
+class _PyTreeExtensionContext:


Not for this PR, this almost feels like deserving a deeper discussion.

If Dynamo cannot support custom types, maybe we could propose an interface for users to register any custom data type, so that exporter knows how to serialize? We can do that at dynamo or exporter context, probably? In fact, we discussed this idea in the past for the torchscript exporter.

thiagocrepaldi · 2023-04-14T19:25:48Z

torch/onnx/_internal/fx/fx_exporter.py

+        self._input_adapter.append_step(adapt_step)
+        return adapt_step.adapt(model_args, model_kwargs)
+
+    def _apply_output_adapt_step(


we could move all these adapters to a file of its own and declutter the exporter class files.

Sure for follow ups

torch/onnx/_internal/exporter.py

…witch to 'DynamoExporter'" Summary * Introduce input/output adapter. Due to design differences, input/output format between PyTorch model and exported ONNX model are often not the same. E.g., `None` inputs are allowed for PyTorch model, but are not supported by ONNX. Nested constructs of tensors are allowed for PyTorch model, but only flattened tensors are supported by ONNX, etc. The new input/output adapter is exported with the model. Providing an interface to automatically convert and validate inputs/outputs format. * As suggested by #98251, provide extension for unwrapping user defined python classes for `dynamo.export` based exporter. Unblock huggingface models. * Re-wire tests to run through `DynamoExporter` w/ `dynamo_export` api. Kept `DynamoOptimizeExporter` in the tests for now for coverage of this change. [ghstack-poisoned]

ghstack-source-id: 0ec4f87e1b5ddb0f774b3da90c467377a42a9a68 Pull Request resolved: #98421

…porter'" Summary * Introduce input/output adapter. Due to design differences, input/output format between PyTorch model and exported ONNX model are often not the same. E.g., `None` inputs are allowed for PyTorch model, but are not supported by ONNX. Nested constructs of tensors are allowed for PyTorch model, but only flattened tensors are supported by ONNX, etc. The new input/output adapter is exported with the model. Providing an interface to automatically convert and validate inputs/outputs format. * As suggested by #98251, provide extension for unwrapping user defined python classes for `dynamo.export` based exporter. Unblock huggingface models. * Re-wire tests to run through `DynamoExporter` w/ `dynamo_export` api. Kept `DynamoOptimizeExporter` in the tests for now for coverage of this change. [ghstack-poisoned]

…amo API usage in fx exporter" Before this PR, dynamo API is not passed in the options.dynamic_shapes leading to a potential static capture in the very begining. After this PR (from dynamo.optimize): ``` def forward(self, L_x_ : torch.Tensor): l_x_ = L_x_ # File: /home/titaiwang/pytorch/test/onnx/test_fx_to_onnx_with_onnxruntime.py:442, code: results.append(x[: x.size(0) - i, i : x.size(2), i:3]) size = l_x_.size(0) sub = size - 0; size = None size_1 = l_x_.size(2) getitem = l_x_[(slice(None, sub, None), slice(0, size_1, None), slice(0, 3, None))]; sub = size_1 = None size_2 = l_x_.size(0) sub_1 = size_2 - 1; size_2 = None size_3 = l_x_.size(2) getitem_1 = l_x_[(slice(None, sub_1, None), slice(1, size_3, None), slice(1, 3, None))]; sub_1 = size_3 = None size_4 = l_x_.size(0) sub_2 = size_4 - 2; size_4 = None size_5 = l_x_.size(2) getitem_2 = l_x_[(slice(None, sub_2, None), slice(2, size_5, None), slice(2, 3, None))]; sub_2 = size_5 = None size_6 = l_x_.size(0) sub_3 = size_6 - 3; size_6 = None size_7 = l_x_.size(2) getitem_3 = l_x_[(slice(None, sub_3, None), slice(3, size_7, None), slice(3, 3, None))]; l_x_ = sub_3 = size_7 = None return (getitem, getitem_1, getitem_2, getitem_3) ``` Wait #98421 to test tracing_mode="symbolic" in `DynamoExportExporter` [ghstack-poisoned]

… fx exporter" Before this PR, dynamo API is not passed in the options.dynamic_shapes leading to a potential static capture in the very begining. After this PR (from dynamo.optimize): ``` def forward(self, L_x_ : torch.Tensor): l_x_ = L_x_ # File: /home/titaiwang/pytorch/test/onnx/test_fx_to_onnx_with_onnxruntime.py:442, code: results.append(x[: x.size(0) - i, i : x.size(2), i:3]) size = l_x_.size(0) sub = size - 0; size = None size_1 = l_x_.size(2) getitem = l_x_[(slice(None, sub, None), slice(0, size_1, None), slice(0, 3, None))]; sub = size_1 = None size_2 = l_x_.size(0) sub_1 = size_2 - 1; size_2 = None size_3 = l_x_.size(2) getitem_1 = l_x_[(slice(None, sub_1, None), slice(1, size_3, None), slice(1, 3, None))]; sub_1 = size_3 = None size_4 = l_x_.size(0) sub_2 = size_4 - 2; size_4 = None size_5 = l_x_.size(2) getitem_2 = l_x_[(slice(None, sub_2, None), slice(2, size_5, None), slice(2, 3, None))]; sub_2 = size_5 = None size_6 = l_x_.size(0) sub_3 = size_6 - 3; size_6 = None size_7 = l_x_.size(2) getitem_3 = l_x_[(slice(None, sub_3, None), slice(3, size_7, None), slice(3, 3, None))]; l_x_ = sub_3 = size_7 = None return (getitem, getitem_1, getitem_2, getitem_3) ``` Wait #98421 to test tracing_mode="symbolic" in `DynamoExportExporter` [ghstack-poisoned]

…8421) Summary * Introduce input/output adapter. Due to design differences, input/output format between PyTorch model and exported ONNX model are often not the same. E.g., `None` inputs are allowed for PyTorch model, but are not supported by ONNX. Nested constructs of tensors are allowed for PyTorch model, but only flattened tensors are supported by ONNX, etc. The new input/output adapter is exported with the model. Providing an interface to automatically convert and validate inputs/outputs format. * As suggested by #98251, provide extension for unwrapping user defined python classes for `dynamo.export` based exporter. Unblock huggingface models. * Re-wire tests to run through `DynamoExporter` w/ `dynamo_export` api. Kept `DynamoOptimizeExporter` in the tests for now for coverage of this change. Pull Request resolved: #98421 Approved by: https://github.com/justinchuby, https://github.com/titaiwangms, https://github.com/thiagocrepaldi

This PR refactors how InputAdapter and OutputAdapter is used throughout the exporter. During refactoring, API issues with passes (torch.onnx._internal.fx._pass.Transform) were identified and should be tackled on another API. In short, some passes can modify the input/output of the model and the input/output adapter must be in sync with such change, otherwise, the adapters will not reflect the actual model input/output. The first instance of this issue was with `ReplaceGetAttrWithPlaceholder` pass that adds new inputs to the model. In order to work this around, a new input adapt step to append new inputs (generated by the pass) was introduced. That resulted in the number of inputs of the ONNX model to mismatch the numer of inputs of the pytorch model, though. Follow up on #98421 Pull Request resolved: #100490 Approved by: https://github.com/BowenBao

This PR refactors how InputAdapter and OutputAdapter is used throughout the exporter. During refactoring, API issues with passes (torch.onnx._internal.fx._pass.Transform) were identified and should be tackled on another API. In short, some passes can modify the input/output of the model and the input/output adapter must be in sync with such change, otherwise, the adapters will not reflect the actual model input/output. The first instance of this issue was with `ReplaceGetAttrWithPlaceholder` pass that adds new inputs to the model. In order to work this around, a new input adapt step to append new inputs (generated by the pass) was introduced. That resulted in the number of inputs of the ONNX model to mismatch the numer of inputs of the pytorch model, though. Follow up on pytorch#98421 Pull Request resolved: pytorch#100490 Approved by: https://github.com/BowenBao

[ONNX] Introduce Input/Ouptut formatter; Switch to 'DynamoExporter'

b331513

[ghstack-poisoned]

pytorch-bot bot added the release notes: onnx torch.onnx related changes that should show up in the release notes label Apr 5, 2023

BowenBao added a commit that referenced this pull request Apr 5, 2023

[ONNX] Introduce Input/Ouptut formatter; Switch to 'DynamoExporter'

4a4b4e3

ghstack-source-id: 39cc6e28980ebd06e249ec3c37e55330434dd2b6 Pull Request resolved: #98421

pytorchbot added the open source label Apr 5, 2023

BowenBao added module: onnx Related to torch.onnx topic: new features topic category labels Apr 5, 2023