[ONNX] Fix type annotations and enable type checking for all apis #84091

justinchuby · 2022-08-25T22:47:07Z

Stack from ghstack (oldest at bottom):

-> [ONNX] Fix type annotations and enable type checking for all apis #84091

Enable runtime type checking for all torch.onnx public apis, symbolic functions and most helpers (minus two that does not have a checkable type: _.JitType does not exist) by adding the beartype decorator. Fix type annotations to makes unit tests green.

Profile:

export torchvision.models.alexnet(pretrained=True)

with runtime type checking: 21.314 / 10 passes
without runtime type checking: 20.797 / 10 passes

+ 2.48%

[ghstack-poisoned]

facebook-github-bot · 2022-08-25T22:47:14Z

🔗 Helpful links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/84091
✖️ Python docs build was skipped
✖️ C++ docs build was skipped
❓Need help or want to give feedback on the CI? Visit our office hours

✅ No Failures (0 Pending)

As of commit 54a44a0 (more details on the Dr. CI page):

Expand to see more

💚 💚 Looks good so far! There are no failures yet. 💚 💚

This comment was automatically generated by Dr. CI (expand for details).

Please report bugs/suggestions to the (internal) Dr. CI Users group.

Click here to manually regenerate this comment.

[ghstack-poisoned]

justinchuby · 2022-08-26T00:33:14Z

torch/onnx/symbolic_opset12.py

-@symbolic_helper.parse_args("v", "f", "i")
+@symbolic_helper.parse_args("v", "f", "b")
+@_beartype.beartype
 def dropout(g, input, p, train):


Changed many "i" to "b" because (1) they are bool types (2) when translated to onnx Booleans are casted to int w/o issues.

justinchuby · 2022-08-26T00:33:42Z

torch/onnx/verification.py

-    ort_outs: Sequence[np.ndarray],
-    pt_outs: Sequence[torch.Tensor],
+    ort_outs: Union[_NumericType, Sequence[_NumericType], Sequence, Dict],
+    pt_outs: Union[_NumericType, Sequence[_NumericType], Sequence, Dict],


Relaxed types here

If ORT can return dict, it means onnx also can. Is that really the case? AFAIK, that was a known limitation, wasn't it? ORTModule flattens dict, export to ONNX, runs the model and reassemble the dict before returning it to the user. Is that the case for allowing dict on ORT? My recolelction was that we used InferenceSession directly, without ORTModule wrapping it to flatten input/output, so it is unexpected to me

Done. Removed Dict

justinchuby · 2022-08-26T00:33:59Z

torch/onnx/verification.py

    use_external_data: bool = False,
-    additional_test_inputs: Optional[Sequence[Tuple[Any, ...]]] = None,
+    additional_test_inputs: Optional[
+        Sequence[Union[torch.Tensor, Tuple[Any, ...]]]


torch/onnx/symbolic_opset9.py

torch/onnx/symbolic_opset11.py

justinchuby · 2022-08-26T00:35:28Z

torch/onnx/symbolic_helper.py

    # NOTE: prim::Constant at this stage usually means something not compatible in ONNX,
    # otherwise it'd be converted to onnx::Constant
-    if _is_value(value) and _is_onnx_constant(value):
+    if isinstance(value, _C.Value) and _is_onnx_constant(value):


Changed for mypy checks on _is_onnx_constant

At least add a TODO comment here and on _is_value describing the mypy issue for a future fix. Otherwise this will be forgotten

justinchuby · 2022-08-26T00:35:40Z

torch/onnx/symbolic_helper.py

-def _maybe_get_const(value: _C.Value, descriptor: _ValueDescriptor):
+@_beartype.beartype
+def _maybe_get_const(
+    value: Optional[Union[_C.Value, torch.Tensor, Number, Sequence]],


Generally speaking, hardening + refactoring is preferred over relaxing checks, especially when the goal of thechange is to enforce correctness

What is the reason for relaxing as opposed to refactor code and removal the avoid the relaxing?

I wanted to (1)enable typing checking before (2) fixing the code. Modifying the annotations is a way to minimal change code to enable checking without modifying current behaviors. We can then optimize stuff from there

justinchuby · 2022-08-26T00:36:15Z

torch/onnx/_patch_torch.py

    g: torch._C.Graph,
    opname: str,
-    *raw_args: torch._C.Value,
+    *raw_args: Union[torch.Tensor, torch._C.Value],


Is this right? Can raw_args be torch.Tensor? (They exist in tests)

I would say no, but could be wrong.
Which test has it? let's look into it

It is expected as I realized there is a const_if_tensor a few lines down

Can we clean up and remove it (preferably in a separate PR)? I guess Meta internal test is likely to scream though :(

Sure thing!

Letting a non Value goes true the most generic function for symbolics open the door for many issues. IMO we need to fix this as opposed to relax the type checking to allow CI to pass. Hardening the code is always best, even if it needs more refactoring

In a follow-up PR. This is already too big

…ll apis" Enable runtime type checking for all torch.onnx public apis, symbolic functions and most helpers (minus two that does not have a checkable type: `_.JitType` does not exist) by adding the beartype decorator. Fix type annotations to makes unit tests green. TODO: profile performance [ghstack-poisoned]

thiagocrepaldi · 2022-08-26T14:54:53Z

Stack from ghstack (oldest at bottom):

-> [ONNX] Fix type annotations and enable type checking for all apis #84091

Enable runtime type checking for all torch.onnx public apis, symbolic functions and most helpers (minus two that does not have a checkable type: _.JitType does not exist) by adding the beartype decorator. Fix type annotations to makes unit tests green.

TODO: profile performance

Add a comment in the places we cannot type check so that nobody tries or do it wrong

Also adds a note stating it is mandatory to add type checking in all symbolics and their helpers are symbolic_helper.py, near the header doc

# Note [Edit Symbolic Files]
# EDITING THIS FILE AND SYMBOLIC_OPSET<VERSION> FILES? READ THIS FIRST!

…g for all apis (#84091) Test Plan: revert-hammer Differential Revision: D39084854 Original commit changeset: aeb0e1fc4dd6 Original Phabricator Diff: D39084854 fbshipit-source-id: 703629a64b2c89c5a4278cbc6a2917b8810e86de

…ll apis" Enable runtime type checking for all torch.onnx public apis, symbolic functions and most helpers (minus two that does not have a checkable type: `_.JitType` does not exist) by adding the beartype decorator. Fix type annotations to makes unit tests green. Profile: export `torchvision.models.alexnet(pretrained=True)` ``` with runtime type checking: 21.314 / 10 passes without runtime type checking: 20.797 / 10 passes + 2.48% ``` [ghstack-poisoned]

ghstack-source-id: cfdc8df Pull Request resolved: #84091

…ll apis" Enable runtime type checking for all torch.onnx public apis, symbolic functions and most helpers (minus two that does not have a checkable type: `_.JitType` does not exist) by adding the beartype decorator. Fix type annotations to makes unit tests green. Profile: export `torchvision.models.alexnet(pretrained=True)` ``` with runtime type checking: 21.314 / 10 passes without runtime type checking: 20.797 / 10 passes + 2.48% ``` [ghstack-poisoned]

ghstack-source-id: d176a12 Pull Request resolved: #84091

justinchuby · 2022-08-30T20:38:15Z

@thiagocrepaldi PTAL cc @BowenBao

justinchuby · 2022-08-30T20:40:10Z

The relaxed types are there to match the reality. We can optimize the code after checks are turned on. The follow ups are also going to be more focused and easier to review

thiagocrepaldi · 2022-08-30T21:14:08Z

LGTM. I noticed that some new annotations were added, while others were not. Is that for passing certain check failures?

same concern here. I'd rather seeing a feature complete on a single PR than spreading towards multiple ones. If it needs to be reverted or cherry picked in the feature, due to a bug or limitation, it is easy to do in a single pr. It is also easier to read a single PR and learn about the introduced feature

thiagocrepaldi

Would rather seeing PRs being code-complete before merge instead of having many follow ups PRs.

That makes reverting unrelated PRs very hard, especially using pytorchbot. (in order to disable beartype, many PRs related and unrelated would have to be reverted (the ones in between all beartype PRs)

It is also harder to learn about the feature because there are too many PRs for the same task

replacing _is_value by isinstance(tensor, _C.Value) defeats the purposa of having is_value()entirely. At least a #TODO or an issue should be created to track and properly fix it.

justinchuby · 2022-08-30T22:56:53Z

Would rather seeing PRs being code-complete before merge instead of having many follow ups PRs.

That makes reverting unrelated PRs very hard, especially using pytorchbot. (in order to disable beartype, many PRs related and unrelated would have to be reverted (the ones in between all beartype PRs)

I like the idea of making revert easier. In this particular case, the bigger the PR becomes and the more places it touches, the harder it can be reverted. Future PRs would more likely need to be reverted because they will build on top of this PR once merged, based on the extend of files it would touch.

IMO this PR stands on its own in the sense that it enables typing checking without substantial changes to the actual code. If we lump the logic changes as well, it will have a higher chance of breaking things and being reverted.

Separating changes to bite sized self-contained PRs often means a smaller chance of breaking things, so we don't need to revert everything. When we start annotating more functions and refactor, the effort can be distributed and verified individually.

I suggest we keep PRs small and focused. https://google.github.io/eng-practices/review/developer/small-cls.html#cant

It is also harder to learn about the feature because there are too many PRs for the same task

PR stacks have been there for this. Depending on what we need, we can create some type of documentation for the feature and do a good job of linking PRs together.

replacing _is_value by isinstance(tensor, _C.Value) defeats the purposa of having is_value()entirely. At least a #TODO or an issue should be created to track and properly fix it.

Will do. thanks!

…ll apis" Enable runtime type checking for all torch.onnx public apis, symbolic functions and most helpers (minus two that does not have a checkable type: `_.JitType` does not exist) by adding the beartype decorator. Fix type annotations to makes unit tests green. Profile: export `torchvision.models.alexnet(pretrained=True)` ``` with runtime type checking: 21.314 / 10 passes without runtime type checking: 20.797 / 10 passes + 2.48% ``` [ghstack-poisoned]

ghstack-source-id: de15bd9 Pull Request resolved: #84091

justinchuby · 2022-09-02T23:48:11Z

@pytorchbot merge -g

pytorchmergebot · 2022-09-02T23:49:44Z

@pytorchbot successfully started a merge job. Check the current status here.
The merge job was triggered with the green (-g) flag. This means that your change will be merged once all checks on your PR have passed (ETA: 0-4 Hours). If this is not the intended behavior, feel free to use some of the other merge options in the wiki.
Please reach out to the PyTorch DevX Team with feedback or questions!

…4091) (#84091) Summary: Enable runtime type checking for all torch.onnx public apis, symbolic functions and most helpers (minus two that does not have a checkable type: `_.JitType` does not exist) by adding the beartype decorator. Fix type annotations to makes unit tests green. Profile: export `torchvision.models.alexnet(pretrained=True)` ``` with runtime type checking: 21.314 / 10 passes without runtime type checking: 20.797 / 10 passes + 2.48% ``` Pull Request resolved: #84091 Approved by: https://github.com/BowenBao, https://github.com/thiagocrepaldi Test Plan: contbuild & OSS CI, see https://hud.pytorch.org/commit/pytorch/pytorch/388368b6996479f6eca484d4e60a6250b2535dec Reviewed By: mehtanirav Differential Revision: D39277677 fbshipit-source-id: 6836efdd15c3b2479bac68807c65ea7c5609295f

Annotate types

ff0fb76

[ghstack-poisoned]

justinchuby requested review from BowenBao and shubhambhokare1 as code owners August 25, 2022 22:47

facebook-github-bot added the cla signed label Aug 25, 2022

justinchuby changed the title ~~Annotate types~~ [ONNX] Annotate types Aug 25, 2022

justinchuby marked this pull request as draft August 25, 2022 22:48

justinchuby added the module: onnx Related to torch.onnx label Aug 25, 2022

pytorchbot added the open source label Aug 25, 2022

Update on "[ONNX] Annotate types"

4b793c3

[ghstack-poisoned]

justinchuby assigned BowenBao Aug 26, 2022

justinchuby marked this pull request as ready for review August 26, 2022 00:29

justinchuby requested a review from thiagocrepaldi August 26, 2022 00:29

justinchuby commented Aug 26, 2022

View reviewed changes

justinchuby added release notes: onnx torch.onnx related changes that should show up in the release notes topic: improvements topic category labels Aug 26, 2022

justinchuby changed the title ~~[ONNX] Annotate types~~ [ONNX] Annotate types and enable type checking for all apis Aug 26, 2022

justinchuby changed the title ~~[ONNX] Annotate types and enable type checking for all apis~~ [ONNX] Fix type annotations and enable type checking for all apis Aug 26, 2022

justinchuby added a commit that referenced this pull request Aug 30, 2022

Annotate types

fa926b4

ghstack-source-id: cfdc8df Pull Request resolved: #84091

justinchuby added a commit that referenced this pull request Aug 30, 2022

Annotate types

6dbe69b

ghstack-source-id: d176a12 Pull Request resolved: #84091

thiagocrepaldi self-requested a review August 30, 2022 21:16

thiagocrepaldi approved these changes Aug 30, 2022

View reviewed changes

justinchuby mentioned this pull request Sep 1, 2022

[ONNX] Add onnx::LayerNorm support for version 17 #84293

Closed

justinchuby added a commit that referenced this pull request Sep 2, 2022

Annotate types

413f7c4

ghstack-source-id: de15bd9 Pull Request resolved: #84091

pytorchmergebot closed this in 388368b Sep 3, 2022

facebook-github-bot deleted the gh/justinchuby/7/head branch September 6, 2022 14:20

[ONNX] Fix type annotations and enable type checking for all apis #84091

[ONNX] Fix type annotations and enable type checking for all apis #84091

Uh oh!

Conversation

justinchuby commented Aug 25, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

facebook-github-bot commented Aug 25, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful links

✅ No Failures (0 Pending)

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!

justinchuby Aug 30, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

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!

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!

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!

thiagocrepaldi Aug 29, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

thiagocrepaldi commented Aug 26, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

justinchuby commented Aug 30, 2022

Uh oh!

justinchuby commented Aug 30, 2022

Uh oh!

thiagocrepaldi commented Aug 30, 2022

Uh oh!

thiagocrepaldi left a comment

Choose a reason for hiding this comment

Uh oh!

justinchuby commented Aug 30, 2022

Uh oh!

justinchuby commented Sep 2, 2022

Uh oh!

pytorchmergebot commented Sep 2, 2022

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

justinchuby commented Aug 25, 2022 •

edited

Loading

facebook-github-bot commented Aug 25, 2022 •

edited

Loading

justinchuby Aug 30, 2022 •

edited

Loading

thiagocrepaldi Aug 29, 2022 •

edited

Loading

thiagocrepaldi commented Aug 26, 2022 •

edited

Loading