[ONNX] Support from dynamic_shapes to dynamic_axes when torch.onnx.export(fallback=True) is triggered #139532

titaiwangms · 2024-11-01T23:07:07Z

Fixes #139320

Summary:

(1) Add `_rename_dynamic_shapes_with_model_inputs` for dynamic_shapes to play along with input_names

Use model forward signature to rename dynamic_shapes when dynamic_shapes is not nested and dynamic_shapes is directly using the customized name. This solves the issue that torch.export.export expects dynamic_shapes only uses the model input names.
If the dynamic_shapes is nested, we do nothing.

(2) Add `_from_dynamic_shapes_to_dynamic_axes` for fallback

We flatten dynamic_shapes with leaf defined _pytree.tree_leaves()
* If a dynamic_shapes is not nested, and defined in dict. We can use the key as the input_names, since it should be renamed by _rename_dynamic_shapes_with_model_inputs already.
If a dynamic_shapes is provided, input_names is required to assign the names, because dynamic_axes needs it.

pytorch-bot · 2024-11-01T23:07:11Z

🔗 Helpful Links

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

📄 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.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

[DomainsOnly] Jobs fail with GLIBC version not found

✅ No Failures

As of commit bdfe02d with merge base e429a3b ():
💚 Looks good so far! There are no failures yet. 💚

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

torch/onnx/_internal/exporter/_compat.py

justinchuby · 2024-11-05T18:54:53Z

torch/onnx/_internal/exporter/_compat.py

+        # It doesn not specify input names if it's a tuple
+        return dynamic_shapes
+
+    sig = _signature(model)


Will this call ever raise?

Yes, when the dynamic_shapes is dict.

Technically dynamic_shapes and args, kwargs are exactly the same (every tensor is replaced by a dictionary of dynamic dimensions) but only a subset of it is supported by torchscript. So i would expect the function to fail if the inputs are not flat unless the module is wrapped into a module doing that. But i don't think we should do it for the users before calling the fallback. However an error message suggesting to do so would be great.

Failing makes sense, and it's easier to implement. I have tried the following example:

import torch import onnx class Foo(torch.nn.Module): def forward(self, x): (a0, a1), (b0, b1), (c0, c1, c2) = x return a0 + a1 + b0 + b1 + c0 + c1 + c2 f = Foo() inputs = ( (1, 2), ( torch.randn(4, 4), torch.randn(4, 4), ), ( torch.randn(4, 4), torch.randn(4, 4), torch.randn(4, 4), ), ) input_names = ["a", "b", "c", "d", "e", "f", "g"] dynamic_axes = { "c": {0: "c_dim_0", 1: "c_dim_1"}, "e": {0: "e_dim_0", 1: "e_dim_1"}, "f": {0: "f_dim_0", 1: "f_dim_1"}, } torch.onnx.export(f, (inputs,), "nested.onnx", dynamic_axes=dynamic_axes, input_names=input_names, verbose=True) onnx_model = onnx.load("nested.onnx") print(onnx_model.graph.input)

I think this is the only way that nested dynamic_axes works, but it's also kind of awkward that model forward input is not the same as in input_names. Users need to flatten ahead in dynamic_axes.

The function has been changed to only work when input is not nested (len(onnx inputs) == len(torch inputs)). When it's nested, users are expected to provide the correct names (match model.forward) of dynamic_shapes, or tuple (honestly, I think users would use tuple when the inputs are nested.)

Could you add this to the function docstring? Just something that makes it clear when this functions works and when it doesn't

torch/onnx/_internal/exporter/_compat.py

…_axes

titaiwangms · 2024-11-06T01:34:50Z

Overall, I think there are a few questions that we need to clarify before adding flattened dynamic_shapes support.

(1) When args is nested, what is our expectation on input_names. for example:

def forward(self, input: tuple[torch.tensor, torch.tensor])  # there are two elements in the tuple

# one to one mapping to forward
input_names = ["input"]
# or flattened to map ONNX nodes
input_names = ["input_a", "input_b"]

(2) dynamic_axes requires users to flatten inputs ahead. For example:

import torch
import onnx

class Foo(torch.nn.Module):
    def forward(self, x):
        (a0, a1), (b0, b1), (c0, c1, c2) = x
        return a0 + a1 + b0 + b1 + c0 + c1 + c2

f = Foo()
inputs = (
    (1, 2),
    (
        torch.randn(4, 4),
        torch.randn(4, 4),
    ),
    (
        torch.randn(4, 4),
        torch.randn(4, 4),
        torch.randn(4, 4),
    ),
)

input_names = ["a", "b", "c", "d", "e", "f", "g"]
dynamic_axes = {
    "c": {0: "c_dim_0", 1: "c_dim_1"},
    "e": {0: "e_dim_0", 1: "e_dim_1"},
    "f": {0: "f_dim_0", 1: "f_dim_1"},
}

torch.onnx.export(f, (inputs,), "nested.onnx", dynamic_axes=dynamic_axes, input_names=input_names, verbose=True)
onnx_model = onnx.load("nested.onnx")
print(onnx_model.graph.input)

Do we want to do this for users? I guess users could hint us by specifying input_names with flattened inputs.

(3) I suggest that we only support nested tuple/list inputs, and we raise error when we see dict in inputs, because (1) torchscript exporter does not support dictionary if I recall correctly, and (2) that's making conversion even more complicated, given this is just a fallback mechanism.

justinchuby · 2024-11-06T04:54:23Z

torchscript does actually support dict inputs, but I think it just flattens the dictionary (may recall incorrectly)

titaiwangms · 2024-11-06T16:33:43Z

torchscript does actually support dict inputs, but I think it just flattens the dictionary (may recall incorrectly)

Oh I was thinking onnx does not support dict. I will make a change.

…_axes

test/onnx/exporter/test_compat.py

justinchuby

Is it true that when a user renames input, they can either provide the old name (in the forward function) or the new name in dynamic shapes? What if they conflict?

…_axes

titaiwangms · 2024-11-15T05:10:21Z

@pytorchbot merge

pytorchmergebot · 2024-11-15T05:12:12Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

torch/onnx/_internal/exporter/_compat.py

pytorchmergebot · 2024-11-15T05:28:11Z

Merge failed

Reason: New commits were pushed while merging. Please rerun the merge command.

Details for Dev Infra team

Raised by workflow job

titaiwangms · 2024-11-15T05:36:11Z

I found a bug. hold on merge for now.

torch/onnx/_internal/exporter/_compat.py

…_axes

justinchuby

Thanks!

justinchuby · 2024-11-18T19:18:25Z

torch/onnx/_internal/exporter/_compat.py

+    *,
+    dynamic_shapes: dict[str, Any] | tuple[Any] | list[Any],
+    input_names: Sequence[str],
+) -> dict[str, Any] | tuple[Any] | list[Any]:


Could you add a docstring to clarify when this function will or will not work?

titaiwangms · 2024-11-18T19:43:08Z

@pytorchbot merge

pytorchmergebot · 2024-11-18T19:45:01Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

…port(fallback=True) is triggered (pytorch#139532) Fixes pytorch#139320 ### Summary: #### (1) Add `_rename_dynamic_shapes_with_model_inputs` for dynamic_shapes to play along with input_names * Use model forward signature to rename dynamic_shapes when dynamic_shapes is not nested and dynamic_shapes is directly using the customized name. This solves the issue that torch.export.export expects dynamic_shapes only uses the model input names. * If the dynamic_shapes is nested, we do nothing. #### (2) Add `_from_dynamic_shapes_to_dynamic_axes` for fallback * We flatten dynamic_shapes with leaf defined _pytree.tree_leaves() ~~* If a dynamic_shapes is not nested, and defined in dict. We can use the key as the input_names, since it should be renamed by `_rename_dynamic_shapes_with_model_inputs` already.~~ * If a dynamic_shapes is provided, input_names is required to assign the names, because dynamic_axes needs it. Pull Request resolved: pytorch#139532 Approved by: https://github.com/justinchuby

titaiwangms added 2 commits November 1, 2024 19:54

first version

decc9ad

add tests

71c055c

titaiwangms requested review from justinchuby, shubhambhokare1 and wschin as code owners November 1, 2024 23:07

pytorch-bot bot added the release notes: onnx torch.onnx related changes that should show up in the release notes label Nov 1, 2024

titaiwangms added the topic: bug fixes topic category label Nov 1, 2024

titaiwangms requested a review from xadupre November 1, 2024 23:13

pytorchbot added the open source label Nov 1, 2024

justinchuby self-assigned this Nov 1, 2024

justinchuby reviewed Nov 5, 2024

View reviewed changes

torch/onnx/_internal/exporter/_compat.py Outdated Show resolved Hide resolved

justinchuby reviewed Nov 5, 2024

View reviewed changes

torch/onnx/_internal/exporter/_compat.py Outdated Show resolved Hide resolved

justinchuby reviewed Nov 5, 2024

View reviewed changes

torch/onnx/_internal/exporter/_compat.py Outdated Show resolved Hide resolved

justinchuby reviewed Nov 5, 2024

View reviewed changes

torch/onnx/_internal/exporter/_compat.py Outdated Show resolved Hide resolved

Merge branch 'main' into titaiwang/add_from_dynamic_shapes_to_dynamic…

3fdc557

…_axes

zou3519 added the triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module label Nov 5, 2024

titaiwangms added 4 commits November 6, 2024 23:34

add support to nested inputs

530f60a

Merge branch 'main' into titaiwang/add_from_dynamic_shapes_to_dynamic…

8716d41

…_axes

delete unused param: model

18d45bf

add export test

a85d65c

titaiwangms requested a review from justinchuby November 8, 2024 23:07

Merge branch 'main' into titaiwang/add_from_dynamic_shapes_to_dynamic…

04243b0

…_axes

justinchuby reviewed Nov 14, 2024

View reviewed changes

test/onnx/exporter/test_compat.py Show resolved Hide resolved

justinchuby approved these changes Nov 14, 2024

View reviewed changes

Merge branch 'main' into titaiwang/add_from_dynamic_shapes_to_dynamic…

b0768b6

…_axes

titaiwangms added 2 commits November 15, 2024 05:04

solve conflict

22fff5a

typo

8aac6e9

pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Nov 15, 2024

pytorchmergebot added the merging label Nov 15, 2024

justinchuby reviewed Nov 15, 2024

View reviewed changes

torch/onnx/_internal/exporter/_compat.py Show resolved Hide resolved

delete dict part

46b664a

pytorchmergebot removed the merging label Nov 15, 2024

justinchuby reviewed Nov 15, 2024

View reviewed changes

torch/onnx/_internal/exporter/_compat.py Outdated Show resolved Hide resolved

justinchuby reviewed Nov 15, 2024

View reviewed changes

torch/onnx/_internal/exporter/_compat.py Outdated Show resolved Hide resolved

titaiwangms added 3 commits November 15, 2024 17:57

add description and minor fix

cb35488

Merge branch 'main' into titaiwang/add_from_dynamic_shapes_to_dynamic…

6c1b7d4

…_axes

add tests

3dc289b

justinchuby approved these changes Nov 18, 2024

View reviewed changes

justinchuby reviewed Nov 18, 2024

View reviewed changes

add docstring

bdfe02d

pytorchmergebot added the merging label Nov 18, 2024

pytorchmergebot added the Merged label Nov 18, 2024

pytorchmergebot closed this in 8e43902 Nov 18, 2024

pytorchmergebot removed the merging label Nov 18, 2024

[ONNX] Support from dynamic_shapes to dynamic_axes when torch.onnx.export(fallback=True) is triggered #139532

[ONNX] Support from dynamic_shapes to dynamic_axes when torch.onnx.export(fallback=True) is triggered #139532

Uh oh!

Conversation

titaiwangms commented Nov 1, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary:

(1) Add _rename_dynamic_shapes_with_model_inputs for dynamic_shapes to play along with input_names

(2) Add _from_dynamic_shapes_to_dynamic_axes for fallback

Uh oh!

pytorch-bot bot commented Nov 1, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

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

❗ 1 Active SEVs

✅ No Failures

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

titaiwangms commented Nov 6, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

justinchuby commented Nov 6, 2024

Uh oh!

titaiwangms commented Nov 6, 2024

Uh oh!

Uh oh!

justinchuby left a comment

Choose a reason for hiding this comment

Uh oh!

titaiwangms commented Nov 15, 2024

Uh oh!

pytorchmergebot commented Nov 15, 2024

Merge started

Uh oh!

Uh oh!

pytorchmergebot commented Nov 15, 2024

Merge failed

Uh oh!

titaiwangms commented Nov 15, 2024

Uh oh!

Uh oh!

Uh oh!

justinchuby left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

titaiwangms commented Nov 18, 2024

Uh oh!

pytorchmergebot commented Nov 18, 2024

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

titaiwangms commented Nov 1, 2024 •

edited

Loading

(1) Add `_rename_dynamic_shapes_with_model_inputs` for dynamic_shapes to play along with input_names

(2) Add `_from_dynamic_shapes_to_dynamic_axes` for fallback

pytorch-bot bot commented Nov 1, 2024 •

edited

Loading

titaiwangms commented Nov 6, 2024 •

edited

Loading