Commit Graph

346 Commits

Author SHA1 Message Date
Thiago Crepaldi
32a67e42c4 Introduce FXGraphExtractor into torch.onnx.dynamo_export (#99940)
The current API architecture can be seen as 3 independent exporters as shown below. The public API `dynamo_export()` defaults to one of the 3 variants and the other 2 must be used by instantiating private classes: ![image](https://user-images.githubusercontent.com/5469809/231567368-ec899718-b7c1-4e59-b6a8-383142df245a.png)

This PR refactors the API in a way that `dynamo_export` is the only way to use the ONNX exporter. It defaults to a FX tracer based on ``torch.export``, but an internal-only idiom allows switching the FX tracer (aka `FXGraphExtractor` interface), as shown below:

![image](https://user-images.githubusercontent.com/5469809/231567495-3936362d-06de-4cfc-b752-6c2060701c08.png)

Summary of changes:

* Unifies all exporter variants under a single `dynamo_export` API
  * `ResolvedExportOptions` was expanded to allow `fx_tracer: FXGraphExtractor` to be specified, selecting which FX graph extractor to use, according to the design proposal
  * As a consequence, `torch.onnx._internal.exporter.Exporter` does not have to *internally* specialize for each type of FX API that the exporter might be used. This leads to a single `Exporter` with many `FX graph extractors`
  * Before in red, after in green: ![image](https://user-images.githubusercontent.com/5469809/232633531-4c67449b-4863-474d-9e18-78fc1d31b1bd.png)
* Input processing was moved from `Exporter` subclasses to `FXGraphExtractor` subclasses, where they are actually consumed
  * `Exporter` is a [data]class that holds export options, model and input data in a single cohesive object. Specializing it means create different exporters instead of having one exporter capable of exporting models through different options.
  * `Exporter` doesn't consume the `model_args` that caused it to specialize
* Improved the circular dependency story.
  * https://github.com/pytorch/pytorch/pull/99070 moves `import torch.onnx` to after all dynamo subcomponents, preventing `torch.onnx` to have circular depemndencies when `torch.XXXX` is imported during initialization
  * There are other points we need to improve in subsequent PRs. APIs are organized in a way that it is easy to "import too much"
* Refactored `decomposition_table` as an internal-only `ResolvedExportOptions` property.
  * Similar to input processing, this helper is not actually consumed at tyhe `Exporter` layer. This PR moves it to the layer in which it is used
* Demoted `Exporter.model_signature` to a simple standalone helper
  * There is no need to have this as a exporter method; this is a standard `inpect.signature` usage without any state

Possible next steps are:
* Decouple `passes` and `dispatching` from the cluttered `export_fx_to_onnx`
* Further integration with http://github.com/pytorch/pytorch/pull/98421/ into `FXGraphExtractor` public API + helper for unit testing
  * Some passes are changing input processing, which are not captured by the proposed input adapter

** COPILOT SUMMARY**
<!--
copilot:all
-->
### <samp>🤖 Generated by Copilot at bdaba31</samp>

### Summary
📝🚀🔧

<!--
1.  📝 - This emoji represents the formatting and documentation changes, such as adding an empty line, updating the `__all__` list, and improving the type annotations and docstrings.
2.  🚀 - This emoji represents the new features and enhancements, such as adding the `DynamoExport` class, supporting custom export options, and flattening HuggingFace model outputs.
3.  🔧 - This emoji represents the refactoring and restructuring changes, such as using the FX graph representation, the `io_adapter` module, and the simplified FX symbolic tracer, and renaming and reorganizing some modules and classes.
-->
This pull request refactors the ONNX exporter code to use the FX graph representation and the new `io_adapter` module for input and output adaptation. It also adds support for custom export options and flattening HuggingFace model outputs in the ONNX test framework. It updates the ONNX dynamo exporter API tests and adds a new module `torch/onnx/_internal/fx/dynamo_graph_extractor.py` for exporting FX models to ONNX with dynamo support. It fixes some type annotations, imports, and formatting issues in the ONNX exporter code.

> _The ONNX exporter got a new look_
> _With FX graph and dynamo hook_
> _It uses `io_adapter`_
> _And custom options matter_
> _For HuggingFace models and `model_signature` book_

### Walkthrough
*  Move the `fx` submodule from `torch/onnx/_internal` to `torch/onnx/_internal/fx`, and rename some of its modules ( [link](https://github.com/pytorch/pytorch/pull/99940/files?diff=unified&w=0#diff-c8fa56eefd7f98fb4f9739d57df57f02ede77e28528133736010a6d06651ebcbL21-R26), [link](https://github.com/pytorch/pytorch/pull/99940/files?diff=unified&w=0#diff-0795f54fd1f38cfbf2c4a863a4efc9f40f2ea020a2b1612605c361b8d8d35862L25-R26), [link](https://github.com/pytorch/pytorch/pull/99940/files?diff=unified&w=0#diff-3eef404cb9d85216c050be153c33255ebce1170a77d8b9b17be79bcfb238c9c4L5-R15), [link](https://github.com/pytorch/pytorch/pull/99940/files?diff=unified&w=0#diff-4da17ba9e1a187bfacb65a70d6ff15f6c2a60480be8e20fc452d8984a279cd0aL3-R30))
*  Add a new module `torch/onnx/_internal/fx/dynamo_graph_extractor.py` that defines a `DynamoExport` class for generating FX graphs using the `torch._dynamo.export` API ([link](https://github.com/pytorch/pytorch/pull/99940/files?diff=unified&w=0#diff-078d7b8d0e4050e650fc3c15dc97a0564852191ac7b7bdc069d0b3959c5ee39aR1-R77))
*  Add a new module `torch/onnx/_internal/fx/io_adapter.py` that defines the input and output adapter classes and steps for the ONNX exporter, and a helper function to wrap models with output adapters ([link](https://github.com/pytorch/pytorch/pull/99940/files?diff=unified&w=0#diff-0795f54fd1f38cfbf2c4a863a4efc9f40f2ea020a2b1612605c361b8d8d35862L159-R192), [link](https://github.com/pytorch/pytorch/pull/99940/files?diff=unified&w=0#diff-4da17ba9e1a187bfacb65a70d6ff15f6c2a60480be8e20fc452d8984a279cd0aL3-R30), [link](https://github.com/pytorch/pytorch/pull/99940/files?diff=unified&w=0#diff-4da17ba9e1a187bfacb65a70d6ff15f6c2a60480be8e20fc452d8984a279cd0aR72-R176), [link](https://github.com/pytorch/pytorch/pull/99940/files?diff=unified&w=0#diff-4da17ba9e1a187bfacb65a70d6ff15f6c2a60480be8e20fc452d8984a279cd0aL237-R478))
*  Update the `ResolvedExportOptions` class in `torch/onnx/_internal/exporter.py` to inherit from the `ExportOptions` class, and to set the `fx_tracer` and `decomposition_table` attributes based on the `dynamo_graph_extractor` and `function_dispatcher` modules ([link](https://github.com/pytorch/pytorch/pull/99940/files?diff=unified&w=0#diff-0795f54fd1f38cfbf2c4a863a4efc9f40f2ea020a2b1612605c361b8d8d35862L81-R99), [link](https://github.com/pytorch/pytorch/pull/99940/files?diff=unified&w=0#diff-0795f54fd1f38cfbf2c4a863a4efc9f40f2ea020a2b1612605c361b8d8d35862R117-R126))
*  Update the `Exporter` class in `torch/onnx/_internal/exporter.py` to remove the `export` method and add a new abstract `generate_fx` method, and to use the `fx_tracer` attribute to generate and export the FX graph ([link](https://github.com/pytorch/pytorch/pull/99940/files?diff=unified&w=0#diff-0795f54fd1f38cfbf2c4a863a4efc9f40f2ea020a2b1612605c361b8d8d35862L413-R475), [link](https://github.com/pytorch/pytorch/pull/99940/files?diff=unified&w=0#diff-0795f54fd1f38cfbf2c4a863a4efc9f40f2ea020a2b1612605c361b8d8d35862L422-R486))
*  Update the `FXSymbolicTraceExporter` class in `torch/onnx/_internal/fx/fx_symbolic_graph_extractor.py` to be renamed to `FXSymbolicTracer`, and to inherit from `exporter.FXGraphExtractor` and implement the `generate_fx` method ([link](https://github.com/pytorch/pytorch/pull/99940/files?diff=unified&w=0#diff-3eef404cb9d85216c050be153c33255ebce1170a77d8b9b17be79bcfb238c9c4L128-R175), [link](https://github.com/pytorch/pytorch/pull/99940/files?diff=unified&w=0#diff-3eef404cb9d85216c050be153c33255ebce1170a77d8b9b17be79bcfb238c9c4L157-R219))
*  Update the `export_fx_to_onnx` method of the `FXSymbolicTracer` class to be renamed to `_export_fx_to_onnx`, and to be moved to the `exporter.FXGraphExtractor` class ([link](https://github.com/pytorch/pytorch/pull/99940/files?diff=unified&w=0#diff-3eef404cb9d85216c050be153c33255ebce1170a77d8b9b17be79bcfb238c9c4L193-R234))
*  Update the `dynamo_export` function in `torch/onnx/_internal/exporter.py` to accept and return `ResolvedExportOptions` and `Exporter` objects, respectively ([link](https://github.com/pytorch/pytorch/pull/99940/files?diff=unified&w=0#diff-0795f54fd1f38cfbf2c4a863a4efc9f40f2ea020a2b1612605c361b8d8d35862L536-R606))
*  Update the `run_test_with_fx_to_onnx_exporter_and_onnx_runtime` function in `test/onnx/onnx_test_common.py` to add a new parameter `export_options` for passing custom export options to the `torch.onnx.dynamo_export` function ([link](https://github.com/pytorch/pytorch/pull/99940/files?diff=unified&w=0#diff-1b38383dc1a0228a835d83bb7c4ba2d0c1bcd41297be5c6336572c525846166eR176), [link](https://github.com/pytorch/pytorch/pull/99940/files?diff=unified&w=0#diff-1b38383dc1a0228a835d83bb7c4ba2d0c1bcd41297be5c6336572c525846166eL216-R222))
*  Update the `test_log_sigmoid` and `_test_large_scale_exporter` tests in `test/onnx/test_fx_to_onnx_with_onnxruntime.py` to use the updated `run_test_with_fx_to_onnx_exporter_and_onnx_runtime` function and the `torch.onnx.dynamo_export` function ([link](https://github.com/pytorch/pytorch/pull/99940/files?diff=unified&w=0#diff-c8fa56eefd7f98fb4f9739d57df57f02ede77e28528133736010a6d06651ebcbL297-R301), [link](https://github.com/pytorch/pytorch/pull/99940/files?diff=unified&w=0#diff-c8fa56eefd7f98fb4f9739d57df57f02ede77e28528133736010a6d06651ebcbL682-R686), [link](https://github.com/pytorch/pytorch/pull/99940/files?diff=unified&w=0#diff-c8fa56eefd7f98fb4f9739d57df57f02ede77e28528133736010a6d06651ebcbL696-R716), [link](https://github.com/pytorch/pytorch/pull/99940/files?diff=unified&w=0#diff-c8fa56eefd7f98fb4f9739d57df57f02ede77e28528133736010a6d06651ebcbL721-R730))
*  Update the `test_raise_on_invalid_save_argument_type` test in `test/onnx/dynamo/test_exporter_api.py` to use the `io_adapter.InputAdapter` and `io_adapter.OutputAdapter` classes instead of the `exporter.InputAdapter` and `exporter.OutputAdapter` classes ([link](https://github.com/pytorch/pytorch/pull/99940/files?diff=unified&w=0#diff-4545f0c15c73ebe90a875e9bee6c5ca4b6b92fb1ed0ec5560d1568e0f6339d02L139-R139))
*  Move the `model_signature` property from the `Exporter` class in `torch/onnx/_internal/exporter.py` to a standalone function in `torch/onnx/utils.py`, and update the references to it ([link](https://github.com/pytorch/pytorch/pull/99940/files?diff=unified&w=0#diff-0795f54fd1f38cfbf2c4a863a4efc9f40f2ea020a2b1612605c361b8d8d35862L432-R505), [link](https://github.com/pytorch/pytorch/pull/99940/files?diff=unified&w=0#diff-3eef404cb9d85216c050be153c33255ebce1170a77d8b9b17be79bcfb238c9c4L157-R219), [link](https://github.com/pytorch/pytorch/pull/99940/files?diff=unified&w=0#diff-849a5778e2dcf7f36587967273cee0bf20642e35bf4c79405111ea3417c3fb3cL54-R75))
*  Move the `UnsatisfiedDependencyError` class from the `Exporter` class in `torch/onnx/_internal/exporter.py` to the top level of the module, and update the references to it ([link](https://github.com/pytorch/pytorch/pull/99940/files?diff=unified&w=0#diff-0795f54fd1f38cfbf2c4a863a4efc9f40f2ea020a2b1612605c361b8d8d35862L442-R512))
*  Rename the `_create_onnx_friendly_decomposition_table` function and the `_ONNX_FRIENDLY_DECOMPOSITION_TABLE` dictionary in `torch/onnx/_internal/fx/function_dispatcher.py` to `_create_default_onnx_decomposition_table` and `_DEFAULT_ONNX_EXPORTER_DECOMPOSITION_TABLE`, respectively, and update the references to them ([link](https://github.com/pytorch/pytorch/pull/99940/files?diff=unified&w=0#diff-549890bc593f917c4e62c4c43077340e4774c0abdf31657ced8450fdfbed3b3eL213-R219), [link](https://github.com/pytorch/pytorch/pull/99940/files?diff=unified&w=0#diff-549890bc593f917c4e62c4c43077340e4774c0abdf31657ced8450fdfbed3b3eL231-R239))
*  Update the imports in `torch/onnx/_internal/fx/function_dispatcher.py` to use the `torch._ops` and `torch._decomp` modules instead of the `torch.ops` and `torch.decomp` modules, and to use aliases for accessing the `onnxscript.function_libs.torch_aten.ops` and `torch._ops` modules ([link](https://github.com/pytorch/pytorch/pull/99940/files?diff=unified&w=0#diff-549890bc593f917c4e62c4c43077340e4774c0abdf31657ced8450fdfbed3b3eL11-R16), [link](https://github.com/pytorch/pytorch/pull/99940/files?diff=unified&w=0#diff-549890bc593f917c4e62c4c43077340e4774c0abdf31657ced8450fdfbed3b3eL35-R156), [link](https://github.com/pytorch/pytorch/pull/99940/files?diff=unified&w=0#diff-549890bc593f917c4e62c4c43077340e4774c0abdf31657ced8450fdfbed3b3eL160-R166), [link](https://github.com/pytorch/pytorch/pull/99940/files?diff=unified&w=0#diff-549890bc593f917c4e62c4c43077340e4774c0abdf31657ced8450fdfbed3b3eL173-R182), [link](https://github.com/pytorch/pytorch/pull/99940/files?diff=unified&w=0#diff-549890bc593f917c4e62c4c43077340e4774c0abdf31657ced8450fdfbed3b3eL189-R194), [link](https://github.com/pytorch/pytorch/pull/99940/files?diff=unified&w=0#diff-549890bc593f917c4e62c4c43077340e4774c0abdf31657ced8450fdfbed3b3eL201-R204), [link](https://github.com/pytorch/pytorch/pull/99940/files?diff=unified&w=0#diff-549890bc593f917c4e62c4c43077340e4774c0abdf31657ced8450fdfbed3b3eL231-R239))
*  Update the `ExportOutput` class in `torch/onnx/_internal/exporter.py` to use the `InputAdapter` and `OutputAdapter` classes from `io_adapter` instead of the ones defined in the same module ([link](https://github.com/pytorch/pytorch/pull/99940/files?diff=unified&w=0#diff-0795f54fd1f38cfbf2c4a863a4efc9f40f2ea020a2b1612605c361b8d8d35862L275-R199))
*  Update the type annotations in `torch/onnx/_internal/fx/serialization.py` and `torch/onnx/_internal/exporter.py` to fix some inconsistencies ([link](https://github.com/pytorch/pytorch/pull/99940/files?diff=unified&w=0#diff-0c7a4333620a22a5c3e5315e30272b59fb7a11b393cb42f8255070bedeb02738L15-R15), [link](https://github.com/pytorch/pytorch/pull/99940/files?diff=unified&w=0#diff-0c7a4333620a22a5c3e5315e30272b59fb7a11b393cb42f8255070bedeb02738L83-R83), [link](https://github.com/pytorch/pytorch/pull/99940/files?diff=unified&w=0#diff-0795f54fd1f38cfbf2c4a863a4efc9f40f2ea020a2b1612605c361b8d8d35862L11-R11), [link](https://github.com/pytorch/pytorch/pull/99940/files?diff=unified&w=0#diff-0795f54fd1f38cfbf2c4a863a4efc9f40f2ea020a2b1612605c361b8d8d35862R18))
*  Remove an unused import of `inspect` from `torch/onnx/_internal/exporter.py` ([link](https://github.com/pytorch/pytorch/pull/99940/files?diff=unified&w=0#diff-0795f54fd1f38cfbf2c4a863a4efc9f40f2ea020a2b1612605c361b8d8d35862L5))
*  Remove an unused import of `torch._dynamo` from `torch/onnx/_internal/fx/passes/shape_inference.py` ([link](https://github.com/pytorch/pytorch/pull/99940/files?diff=unified&w=0#diff-d38827b1f79525963c39e5c480240cd81f4edcaf8b3bd374a1c6ee2fdb28b334L7))
*  Add a comment to `torch/onnx/_internal/fx/passes/shape_inference.py` to explain why the import of `torch._dynamo` is done inside the `_run` method of the `ShapeInferenceWithFakeTensor` class ([link](https://github.com/pytorch/pytorch/pull/99940/files?diff=unified&w=0#diff-d38827b1f79525963c39e5c480240cd81f4edcaf8b3bd374a1c6ee2fdb28b334R32-R35))
*  Fix a typo in the docstring of the `_module_expansion_symbolic_trace` function in `torch/onnx/_internal/fx/fx_symbolic_graph_extractor.py` ([link](https://github.com/pytorch/pytorch/pull/99940/files?diff=unified&w=0#diff-3eef404cb9d85216c050be153c33255ebce1170a77d8b9b17be79bcfb238c9c4L96-R98))
*  Add an empty line to `torch/onnx/__init__.py` for formatting purposes ([link](https://github.com/pytorch/pytorch/pull/99940/files?diff=unified&w=0#diff-c3c8c09b65c1235ca4494633c6a0aab2761a11a7653ddaf9f874bbcd91e15553R12))
*  Delete the `torch/onnx/_internal/fx/__init__.py` file ([link](https://github.com/pytorch/pytorch/pull/99940/files?diff=unified&w=0#diff-a39fa3741f027bb9717388fc922d1e846fbd43d44f2c5fbee4e8d2188a7edb85))

Fixes #ISSUE_NUMBER

Pull Request resolved: https://github.com/pytorch/pytorch/pull/99940
Approved by: https://github.com/BowenBao, https://github.com/jansel
2023-04-27 00:25:28 +00:00
Aaron Gokaslan
dd5e6e8553 [BE]: Merge startswith calls - rule PIE810 (#96754)
Merges startswith, endswith calls to into a single call that feeds in a tuple. Not only are these calls more readable, but it will be more efficient as it iterates through each string only once.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/96754
Approved by: https://github.com/ezyang
2023-03-14 22:05:20 +00:00
Thiago Crepaldi
b9c25f186c Ignore shape inference exception from Caffe2 ATen fallback (#90408)
Fixes #87318

Pull Request resolved: https://github.com/pytorch/pytorch/pull/90408
Approved by: https://github.com/BowenBao
2023-03-08 20:02:11 +00:00
PyTorch MergeBot
3f840cc627 Revert "Ignore shape inference exception from Caffe2 ATen fallback (#90408)"
This reverts commit 1d4e872370.

Reverted https://github.com/pytorch/pytorch/pull/90408 on behalf of https://github.com/huydhn due to Sorry for reverting your PR, but it breaks lint check https://hud.pytorch.org/pr/90408#11855039599. Please fix the error and reland your change
2023-03-08 17:28:21 +00:00
Thiago Crepaldi
1d4e872370 Ignore shape inference exception from Caffe2 ATen fallback (#90408)
Fixes #87318

Pull Request resolved: https://github.com/pytorch/pytorch/pull/90408
Approved by: https://github.com/BowenBao
2023-03-08 16:57:48 +00:00
Thiago Crepaldi
a9a3a1bd14 Apply peephole for eval mode when constant folding is enabled only (#95801)
Fixes https://github.com/microsoft/onnx-converters-private/issues/150

Pull Request resolved: https://github.com/pytorch/pytorch/pull/95801
Approved by: https://github.com/BowenBao
2023-03-01 23:07:38 +00:00
Thiago Crepaldi
d978395f55 Deprecate Caffe2 ONNX exporter (#94994)
Discussed on Weekly meeting with Meta on 2/16/2023 with @kit1980 @malfet

Pull Request resolved: https://github.com/pytorch/pytorch/pull/94994
Approved by: https://github.com/Skylion007, https://github.com/BowenBao
2023-02-17 15:41:11 +00:00
Xuehai Pan
b005ec62b9 [BE] Remove dependency on six and future (#94709)
Remove the Python 2 and 3 compatibility library [six](https://pypi.org/project/six) and [future](https://pypi.org/project/future) and `torch._six`. We only support Python 3.8+ now. It's time to retire them.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/94709
Approved by: https://github.com/malfet, https://github.com/Skylion007
2023-02-14 09:14:14 +00:00
Justin Chu
6cef200af9 [ONNX] Wrap symbolic method calls with graph context (#94746)
This should address #93370

Pull Request resolved: https://github.com/pytorch/pytorch/pull/94746
Approved by: https://github.com/BowenBao
2023-02-13 21:29:39 +00:00
Aaron Gokaslan
9171f7d4cd [BE] Modernize PyTorch even more for 3.8 with pyupgrade (#94520)
Applies some more pyupgrade fixits to PyTorch

Pull Request resolved: https://github.com/pytorch/pytorch/pull/94520
Approved by: https://github.com/ezyang
2023-02-10 18:02:50 +00:00
AllenTiTaiWang
04b06c9627 [ONNX] Use optional op to keep None in results for ONNX internal tests (#84789)
All this time, PyTorch and ONNX has different strategy for None in output. And in internal test, we flatten the torch outputs to see if the rest of them matched. However, this doesn't work anymore in scripting after Optional node is introduced, since some of None would be kept.

#83184 forces script module to keep all Nones from Pytorch, but in ONNX, the model only keeps the ones generated with Optional node, and deletes those meaningless None.

This PR uses Optional node to keep those meaningless None in output as well, so when it comes to script module result comparison, Pytorch and ONNX should have the same amount of Nones.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/84789
Approved by: https://github.com/BowenBao
2023-02-08 23:04:47 +00:00
AllenTiTaiWang
b27ac6dc56 [ONNX] Add full checker mode in torch.onnx.export (#83186)
Fix #82589
Why:
1. **full_check** works in `onnx::checker::check_model` function as it turns on **strict_mode** in `onnx::shape_inference::InferShapes()` which I think that was the intention of this part of code.
2. **strict_mode** catches failed shape type inference (invalid ONNX model from onnx perspective) and ONNXRUNTIME can't run these invalid models, as ONNXRUNTIME actually rely on ONNX shape type inference to optimize ONNX graph. Why we don't set it True for default? >>> some of existing users use other platform, such as caffe2 to run ONNX model which doesn't need valid ONNX model to run.
3. This PR doesn't change the original behavior of `check_onnx_proto`, but add a warning message for those models which can't pass strict shape type inference, saying the models would fail on onnxruntime.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/83186
Approved by: https://github.com/justinchuby, https://github.com/thiagocrepaldi, https://github.com/jcwchen, https://github.com/BowenBao
2023-02-08 22:47:25 +00:00
Ram Rachum
77f9b2e8bf Fix exception causes in fx, nn and onnx packages (#90134)
This is a continuation of #90118

@kit1980
Pull Request resolved: https://github.com/pytorch/pytorch/pull/90134
Approved by: https://github.com/kit1980
2022-12-06 04:34:58 +00:00
Keval Morabia
3d247a8bcd Fix unconvertible_ops as per #89261 (#89299)
Fixes #89261

Pull Request resolved: https://github.com/pytorch/pytorch/pull/89299
Approved by: https://github.com/justinchuby, https://github.com/BowenBao
2022-11-21 20:40:04 +00:00
AllenTiTaiWang
c3acb9c885 [ONNX] Add Internal Utils: onnx_proto_utils.py for onnx/onnx-script/onnx_proto (#88376)
Added `onnx_proto_utils.py` for onnx/onnx-script related process. The idea is like jit_utils.py, and to simplify what we have in `torch/onnx/utils.py`

Pull Request resolved: https://github.com/pytorch/pytorch/pull/88376
Approved by: https://github.com/justinchuby, https://github.com/BowenBao
2022-11-17 03:08:09 +00:00
AllenTiTaiWang
abe41aee77 [ONNX] Support custom Op with onnx-script local function (#86906)
Extend `register_custom_op` to support onnx-script local function. The FunctionProto from onnx-script is represented by custom op and inserted into ModelProto for op execution.

NOTE: I did experiments on >2GB case of a simple model with large initializers:

```python
import torch

class Net(torch.nn.Module):
    def __init__(self, B, C):
        super().__init__()
        self.layer_norm = torch.nn.LayerNorm((B, C), eps=1e-3)
    def forward(self, x):
        return self.layer_norm(x)

N, B, C = 3, 25000, 25000
model = Net(B, C)
x = torch.randn(N, B, C)

torch.onnx.export(model, x, "large_model.onnx", opset_version=12)
```

And it turns out we won't get model_bytes > 2GB after `_export_onnx` pybind cpp function, as we split initializer in external files in that function, and have serialization before return the model bytes, which protobuf is not allowed to be larger than 2GB at any circumstances.

The test cases can be found in the next PR #86907 .

Pull Request resolved: https://github.com/pytorch/pytorch/pull/86906
Approved by: https://github.com/justinchuby, https://github.com/BowenBao
2022-11-16 15:08:55 +00:00
Thiago Crepaldi
5f0783bd6d Fix ATen Fallback for BUILD_CAFFE2=0 for ONNX-only ops (#88504)
Follow-up for #87735

Once again, because BUILD_CAFFE2=0 is not tested for ONNX exporter, one scenario slipped through. A use case where the model can be exported without aten fallback when operator_export_type=ONNX_ATEN_FALLBACK and BUILD_CAFFE2=0

A new unit test has been added, but it won't prevent regressions if BUILD_CAFFE2=0 is not executed on CI again

Fixes #87313

Pull Request resolved: https://github.com/pytorch/pytorch/pull/88504
Approved by: https://github.com/justinchuby, https://github.com/BowenBao
2022-11-11 17:43:46 +00:00
Thiago Crepaldi
2aed670710 Fix ONNX operator_export_type on the new registry (#87735)
Fixes #87313

Our ONNX pipelines do not run with BUILD_CAFFE2=0, so tests for operator_export_type ONNX_ATEN and ONNX_ATEN_FALLBACK will not be fully tested, allowing regressions to happen again.

We need to run the same set of tests for both BUILD_CAFFE2=0 and 1
Pull Request resolved: https://github.com/pytorch/pytorch/pull/87735
Approved by: https://github.com/AllenTiTaiWang, https://github.com/BowenBao
2022-11-02 15:54:40 +00:00
AllenTiTaiWang
4210cebc16 [ONNX] Add internal node kind parsing (#87638)
Pull Request resolved: https://github.com/pytorch/pytorch/pull/87638
Approved by: https://github.com/justinchuby, https://github.com/BowenBao
2022-10-29 11:51:23 +00:00
Justin Chu
c600ce39ed [ONNX] Refactor UnsupportedOperatorError arguments (#85349)
Merged the first two arguments because we always use qualified names to identify symbolic functions
Pull Request resolved: https://github.com/pytorch/pytorch/pull/85349
Approved by: https://github.com/AllenTiTaiWang, https://github.com/BowenBao
2022-10-26 00:21:58 +00:00
Justin Chu
56a744bf47 [ONNX] Reland: Update training state logic to support ScriptedModule (#86745)
In https://github.com/pytorch/pytorch/issues/86325, it was reported that ScriptedModule do not have a training attribute and will fail export because we don't expect them as input.

Also

- Parameterized the test_util_funs test

Thanks @borisfom for the suggestion!

Fixes #86325

Pull Request resolved: https://github.com/pytorch/pytorch/pull/86745
Approved by: https://github.com/AllenTiTaiWang, https://github.com/BowenBao
2022-10-14 19:44:47 +00:00
PyTorch MergeBot
056cfb0464 Revert "[ONNX] Update training state logic to support ScriptedModule (#86745)"
This reverts commit 960b98128e.

Reverted https://github.com/pytorch/pytorch/pull/86745 on behalf of https://github.com/janeyx99 due to  960b98128e broke onnx tests on trunk
2022-10-14 05:40:20 +00:00
Justin Chu
960b98128e [ONNX] Update training state logic to support ScriptedModule (#86745)
In https://github.com/pytorch/pytorch/issues/86325, it was reported that ScriptedModule do not have a training attribute and will fail export because we don't expect them as input.

Also

- Parameterized the test_util_funs test

Thanks @borisfom for the suggestion!

Fixes #86325

Pull Request resolved: https://github.com/pytorch/pytorch/pull/86745
Approved by: https://github.com/AllenTiTaiWang, https://github.com/BowenBao
2022-10-14 01:31:40 +00:00
Justin Chu
69b927701a [ONNX] Update user documentation (#85819)
- Remove mentions of `SymbolicContext` in the doc
- Comment out the PythonOp example so that it is not shown to users
- Updated code blocks and wording
- Changed to recommend using `pip` for installing onnx.

Now adds a deprecation message to the docs (demo only):

![image](https://user-images.githubusercontent.com/11205048/193327649-f789b369-6b59-49e0-8bba-34a6785eb128.png)

Fixes #85608

Pull Request resolved: https://github.com/pytorch/pytorch/pull/85819
Approved by: https://github.com/AllenTiTaiWang, https://github.com/BowenBao
2022-09-30 19:35:34 +00:00
BowenBao
8f4edf1e1d [ONNX] Initial version of diagnostics infrastructure. (#85107)
This PR introduces a general Python diagnostics infrastructure powered by SARIF,
and the exporter diagnostics module that builds on top of it.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/85107
Approved by: https://github.com/abock, https://github.com/justinchuby
2022-09-30 07:47:26 +00:00
Justin Chu
dc63948dc9 [ONNX] Update behavior for register_custom_op_symbolic (#85636)
Update `register_custom_op_symbolic`'s behavior to _only register the symbolic function at a single version_. This is more aligned with the semantics of the API signature.

As a result of this change, opset 7 and opset 8 implementations are now seen as fallback when the opset_version >= 9. Previously any ops internally registered to opset < 9 are not discoverable by an export version target >= 9. Updated the test to reflect this change.

The implication of this change is that users will need to register a symbolic function to the exact version when they want to override an existing symbolic. They are not impacted if (1) an implementation does not existing for the op, or (2) they are already registering to the exact version for export.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/85636
Approved by: https://github.com/BowenBao
2022-09-29 04:24:06 +00:00
Justin Chu
7cdd39b393 [ONNX] Update unconvertible_ops (#85595)
Update `unconvertible_ops` to create a list of unconvertible ops using the updated registry.

- Use fewer passes in the jit process instead to avoid errors during conversion in the ONNX fallback mode
- Actually check the registry to find implemented ops
- Fix type hints for `_create_jit_graph` and `_jit_pass_onnx_remove_inplace_ops_for_onnx`
Pull Request resolved: https://github.com/pytorch/pytorch/pull/85595
Approved by: https://github.com/BowenBao
2022-09-29 00:52:21 +00:00
Justin Chu
85d8441fba [ONNX] Deprecate setter functions for global variables (#85165)
`_set_opset_version` and `_set_operator_export_type` are previously deprecated. This PR decorates them with the deprecation decorator, so warnings are emitted.

- Remove usage of `_set_opset_version` and `_set_operator_export_type` in favor of setting the globals vars directly in torch.onnx internal
- Update `GLOBALS.operator_export_type`'s default to not be None to tighten types
- Remove usage of `_set_onnx_shape_inference`
Pull Request resolved: https://github.com/pytorch/pytorch/pull/85165
Approved by: https://github.com/BowenBao, https://github.com/AllenTiTaiWang
2022-09-28 22:43:43 +00:00
Justin Chu
c42a408baa [ONNX] Create decorator to handle symbolic context (#84776)
- Create decorator to handle old style custom symbolics that require context
- Deprecate `torch.onnx.SymbolicContext` in favor of `GraphContext`. Added deprecation message
- Remove README reference of SymbolicContext

Pull Request resolved: https://github.com/pytorch/pytorch/pull/84776
Approved by: https://github.com/AllenTiTaiWang, https://github.com/BowenBao
2022-09-28 22:36:54 +00:00
Justin Chu
3d2316670f [ONNX] Create GraphContext and load g.op method to the class (#84728)
This PR create the `GraphContext` class and relays all graph methods to _C.Graph as well as implements the `g.op`  method. The GraphContext object is passed into the symbolic functions in place of _C.Graph for compatibility with existing symbolic functions.

This way (1) we can type annotate all `g` args because the method is defined and (2) we can use additional context information in symbolic functions. (3) no more monkey patching on `_C.Graph`

Also

- Fix return type of `_jit_pass_fixup_onnx_controlflow_node`
- Create `torchscript.py` to house torch.Graph related functions
- Change `GraphContext.op` to create nodes in the Block instead of the Graph
- Create `add_op_with_blocks` to handle scenarios where we need to directly manipulate sub-blocks. Update loop and if symbolic functions to use this function.

## Discussion

Should we put all the context inside `SymbolicContext` and make it an attribute in the `GraphContext` class? This way we only define two attributes `GraphContext.graph` and `GraphContext.context`. Currently all context attributes are directly defined in the class.

### Decision

Keep GraphContext flatand note that it will change in the future.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/84728
Approved by: https://github.com/AllenTiTaiWang, https://github.com/BowenBao
2022-09-28 22:21:55 +00:00
Thiago Crepaldi
6cfe555f4f [ONNX] Apply Common Subexpression Elimination pass to ONNX export (#85665)
## Summary
Exporting graphs with Autocast may fail due to a limitation on JIT tracer. By disabling Autocast cache, tracer works, but there can be performance hit when there is reuse of weights in convolution, for example

By applying CSE, such performance loss can be reverted.

ps: As a comment at #84092 mentioned, disabling Autocast cache is an acceptable workaround and used throughout PyTorch code.

Fixes #84092

## Examples of before and after CSE being applied:

### Example: eliminating `%17` and reusing `%16` instead

```python
# BEFORE
graph(%0 : Float(requires_grad=0, device=cpu)):
  %3 : Scalar = aten::ScalarImplicit(%0), scope: test_onnx_opset.TestONNXOpset.test_full.<locals>.MyModule::
  %13 : int = prim::Constant[value=3](), scope: test_onnx_opset.TestONNXOpset.test_full.<locals>.MyModule:: # /home/thiagofc/dev/github/pytorch/test/onnx/test_onnx_opset.py:347:0
  %14 : int = prim::Constant[value=4](), scope: test_onnx_opset.TestONNXOpset.test_full.<locals>.MyModule:: # /home/thiagofc/dev/github/pytorch/test/onnx/test_onnx_opset.py:347:0
  %15 : int[] = prim::ListConstruct(%13, %14), scope: test_onnx_opset.TestONNXOpset.test_full.<locals>.MyModule::
  %16 : NoneType = prim::Constant(), scope: test_onnx_opset.TestONNXOpset.test_full.<locals>.MyModule::
  %17 : NoneType = prim::Constant(), scope: test_onnx_opset.TestONNXOpset.test_full.<locals>.MyModule::
  %18 : Device = prim::Constant[value="cpu"](), scope: test_onnx_opset.TestONNXOpset.test_full.<locals>.MyModule:: # /home/thiagofc/dev/github/pytorch/test/onnx/test_onnx_opset.py:347:0
  %19 : bool = prim::Constant[value=0](), scope: test_onnx_opset.TestONNXOpset.test_full.<locals>.MyModule:: # /home/thiagofc/dev/github/pytorch/test/onnx/test_onnx_opset.py:347:0
  %20 : Float(3, 4, strides=[4, 1], requires_grad=0, device=cpu) = aten::full(%15, %3, %16, %17, %18, %19), scope: test_onnx_opset.TestONNXOpset.test_full.<locals>.MyModule:: # /home/thiagofc/dev/github/pytorch/test/onnx/test_onnx_opset.py:347:0
  return (%20)

AFTER
graph(%0 : Float(requires_grad=0, device=cpu)):
  %3 : Scalar = aten::ScalarImplicit(%0), scope: test_onnx_opset.TestONNXOpset.test_full.<locals>.MyModule::
  %13 : int = prim::Constant[value=3](), scope: test_onnx_opset.TestONNXOpset.test_full.<locals>.MyModule:: # /home/thiagofc/dev/github/pytorch/test/onnx/test_onnx_opset.py:347:0
  %14 : int = prim::Constant[value=4](), scope: test_onnx_opset.TestONNXOpset.test_full.<locals>.MyModule:: # /home/thiagofc/dev/github/pytorch/test/onnx/test_onnx_opset.py:347:0
  %15 : int[] = prim::ListConstruct(%13, %14), scope: test_onnx_opset.TestONNXOpset.test_full.<locals>.MyModule::
  %16 : NoneType = prim::Constant(), scope: test_onnx_opset.TestONNXOpset.test_full.<locals>.MyModule::
  %18 : Device = prim::Constant[value="cpu"](), scope: test_onnx_opset.TestONNXOpset.test_full.<locals>.MyModule:: # /home/thiagofc/dev/github/pytorch/test/onnx/test_onnx_opset.py:347:0
  %19 : bool = prim::Constant[value=0](), scope: test_onnx_opset.TestONNXOpset.test_full.<locals>.MyModule:: # /home/thiagofc/dev/github/pytorch/test/onnx/test_onnx_opset.py:347:0
  %20 : Float(3, 4, strides=[4, 1], requires_grad=0, device=cpu) = aten::full(%15, %3, %16, %16, %18, %19), scope: test_onnx_opset.TestONNXOpset.test_full.<locals>.MyModule:: # /home/thiagofc/dev/github/pytorch/test/onnx/test_onnx_opset.py:347:0
  return (%20)
```
Pull Request resolved: https://github.com/pytorch/pytorch/pull/85665
Approved by: https://github.com/ngimel, https://github.com/AllenTiTaiWang, https://github.com/BowenBao
2022-09-27 21:26:32 +00:00
Justin Chu
2f50d2f685 [ONNX] Update docs on symbolic registration (#85290)
- Move inline instructions on editing symbolic functions to the README
- Add a line on using the symbolic function registration decorator.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/85290
Approved by: https://github.com/BowenBao
2022-09-22 13:37:11 +00:00
Justin Chu
9d1155235b [ONNX] Create decorators for symbolic function registration (#84709)
This PR creates and tests the decorators proposed in #83787

Pull Request resolved: https://github.com/pytorch/pytorch/pull/84709
Approved by: https://github.com/AllenTiTaiWang, https://github.com/BowenBao
2022-09-17 01:01:04 +00:00
Justin Chu
cd7e6d4ad1 [ONNX] New symbolic function registry (#84382)
## Summary

The change brings the new registry for symbolic functions in ONNX. The `SymbolicRegistry` class in `torch.onnx._internal.registration` replaces the dictionary and various functions defined in `torch.onnx.symbolic_registry`.

The new registry

- Has faster lookup by storing only functions in the opset version they are defined in
- Is easier to manage and interact with due to its class design
- Builds the foundation for the more flexible registration process detailed in #83787

Implementation changes

- **Breaking**: Remove `torch.onnx.symbolic_registry`
- `register_custom_op_symbolic` and `unregister_custom_op_symbolic` in utils maintain their api for compatibility
- Update _onnx_supported_ops.py for doc generation to include quantized ops.
- Update code to register python ops in `torch/csrc/jit/passes/onnx.cpp`

## Profiling results

-0.1 seconds in execution time. -34% time spent in `_run_symbolic_function`. Tested on the alexnet example in public doc.

### After
```
   └─ 1.641 export  <@beartype(torch.onnx.utils.export) at 0x7f19be17f790>:1
      └─ 1.641 export  torch/onnx/utils.py:185
         └─ 1.640 _export  torch/onnx/utils.py:1331
            ├─ 0.889 _model_to_graph  torch/onnx/utils.py:1005
            │  ├─ 0.478 _optimize_graph  torch/onnx/utils.py:535
            │  │  ├─ 0.214 PyCapsule._jit_pass_onnx_graph_shape_type_inference  <built-in>:0
            │  │  │     [2 frames hidden]  <built-in>
            │  │  ├─ 0.190 _run_symbolic_function  torch/onnx/utils.py:1670
            │  │  │  └─ 0.145 Constant  torch/onnx/symbolic_opset9.py:5782
            │  │  │     └─ 0.139 _graph_op  torch/onnx/_patch_torch.py:18
            │  │  │        └─ 0.134 PyCapsule._jit_pass_onnx_node_shape_type_inference  <built-in>:0
            │  │  │              [2 frames hidden]  <built-in>
            │  │  └─ 0.033 [self]
```

### Before
![image](https://user-images.githubusercontent.com/11205048/188032302-688d881e-860d-4046-bdba-90da54233576.png)

### Start up time

The startup process takes 0.03 seconds. Calls to `inspect` will be eliminated when we switch to using decorators for registration in #84448

![image](https://user-images.githubusercontent.com/11205048/188208910-250f0434-475d-4872-9abc-781535519305.png)

Pull Request resolved: https://github.com/pytorch/pytorch/pull/84382
Approved by: https://github.com/AllenTiTaiWang, https://github.com/BowenBao
2022-09-16 21:45:16 +00:00
Justin Chu
2fa8142cf9 [ONNX] Rename constants for clarity (#84645)
Rename constants to make them more clear. Fix styles to upper case.

Removed `onnx_stable_opsets` because it can be computed from `ONNX_MIN_OPSET` and `ONNX_MAX_OPSET`.

Fixes #84643

Pull Request resolved: https://github.com/pytorch/pytorch/pull/84645
Approved by: https://github.com/BowenBao
2022-09-09 01:22:14 +00:00
Thibault
d9ceda49c4 ONNX: fix default function value in _optimize_graph (#83996)
The default value for params_dict in _optimize_graph, which is None, throw the following error:

>     _C._jit_pass_onnx_unpack_quantized_weights(
> TypeError: _jit_pass_onnx_unpack_quantized_weights(): incompatible function arguments. The following argument types are supported:
>     1. (arg0: torch::jit::Graph, arg1: Dict[str, IValue], arg2: bool) -> Dict[str, IValue]

Replacing it by an empty dict fixes the issue (and makes more sense).

Pull Request resolved: https://github.com/pytorch/pytorch/pull/83996
Approved by: https://github.com/BowenBao
2022-09-06 23:32:16 +00:00
Justin Chu
388368b699 [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%
```
Pull Request resolved: https://github.com/pytorch/pytorch/pull/84091
Approved by: https://github.com/BowenBao, https://github.com/thiagocrepaldi
2022-09-03 01:40:18 +00:00
titaiwang
ece0002c4b [ONNX] Disable autocast cache in exporter (#84219)
This PR provides a temporary fix on #84092 in exporter to avoid more cases falling into this bug.
A long-term fix will be provided later.

A simple repro with torch.onnx.export is still under investigation, as torch.jit.trace() is not the API we call inside torch.onnx.export, and it may introduce the difference. Therefore, a test case is provided here only.
A specific test one can use,
```python
import torch
import onnxruntime
from onnxruntime.training.ortmodule import DebugOptions, LogLevel
from onnxruntime.training.ortmodule import ORTModule

class MyModule(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.cv1 = torch.nn.Conv2d(3, 3, 5, 2, 1)

    def forward(self, x):
        x = self.cv1(x)
        return x

x = torch.randn(10, 3, 20, 20) * 2
m = MyModule().eval()
x = x.cuda()
m = m.cuda()

debug_options = DebugOptions(log_level=LogLevel.VERBOSE, save_onnx=True, onnx_prefix="ViT-B")
m = ORTModule(m, debug_options=debug_options)

with torch.cuda.amp.autocast(dtype=torch.float16, cache_enabled=True):
    loss = m(x)
```
AND make assertion fail in ORTModule
17ccd6fa02/orttraining/orttraining/python/training/ortmodule/_io.py (L578-L581)

Without the fix, the user will see the weight/bias of Conv node becomes constant.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/84219
Approved by: https://github.com/BowenBao, https://github.com/thiagocrepaldi
2022-09-01 00:34:37 +00:00
titaiwang
18264432f7 [ONNX] replace all _C._flatten to torch.jit._flatten (#83598)
_C._flatten is exactly the same as torch.jit._flatten. Unifying them to reduce confusion.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/83598
Approved by: https://github.com/justinchuby, https://github.com/BowenBao
2022-09-01 00:31:28 +00:00
BowenBao
806878518f [ONNX][Reland] Export node and value with scope name (#82040)
Introduce `_jit_pass_onnx_assign_node_and_value_names` to parse and assign
scoped name for nodes and values in exported onnx graph.
Module layer information is obtained from `ONNXScopeName` captured in `scope`
attribute in nodes. For nodes, the processed onnx node name are stored in
attribute `onnx_name`. For values, the processed onnx output name are stored
as `debugName`.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/82040
Approved by: https://github.com/AllenTiTaiWang, https://github.com/justinchuby, https://github.com/abock
2022-08-29 20:10:38 +00:00
PyTorch MergeBot
8e6207bcd8 Revert "[ONNX] Export node and value with scope name (#82040)"
This reverts commit 6a3666282d.

Reverted https://github.com/pytorch/pytorch/pull/82040 on behalf of https://github.com/weiwangmeta due to Diff reverted internally
2022-08-29 06:36:18 +00:00
PyTorch MergeBot
d8cc8368ab Revert "[ONNX] Fix type annotations and enable type checking for all apis (#84091)"
This reverts commit 6446da1730.

Reverted https://github.com/pytorch/pytorch/pull/84091 on behalf of https://github.com/facebook-github-bot due to Diff reverted internally
2022-08-28 12:28:58 +00:00
Justin Chu
6446da1730 [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%
```
Pull Request resolved: https://github.com/pytorch/pytorch/pull/84091
Approved by: https://github.com/BowenBao
2022-08-27 04:40:41 +00:00
BowenBao
6a3666282d [ONNX] Export node and value with scope name (#82040)
Introduce `_jit_pass_onnx_assign_node_and_value_names` to parse and assign
scoped name for nodes and values in exported onnx graph.
Module layer information is obtained from `ONNXScopeName` captured in `scope`
attribute in nodes. For nodes, the processed onnx node name are stored in
attribute `onnx_name`. For values, the processed onnx output name are stored
as `debugName`.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/82040
Approved by: https://github.com/AllenTiTaiWang, https://github.com/justinchuby, https://github.com/abock
2022-08-26 20:59:12 +00:00
Justin Chu
bf25a140f9 [ONNX] Add runtime type checking to export (#83673)
This PR adds an internal wrapper on the [beartype](https://github.com/beartype/beartype) library to perform runtime type checking in `torch.onnx`. It uses beartype when it is found in the environment and is reduced to a no-op when beartype is not found.

Setting the env var `TORCH_ONNX_EXPERIMENTAL_RUNTIME_TYPE_CHECK=ERRORS` will turn on the feature. setting `TORCH_ONNX_EXPERIMENTAL_RUNTIME_TYPE_CHECK=DISABLED` will disable all checks. When not set and `beartype` is installed, a warning message is emitted.

Now when users call an api with invalid arguments e.g.

```python
torch.onnx.export(conv, y, path, export_params=True, training=False)

# traning should take TrainingModel, not bool
```

they get

```
Traceback (most recent call last):
  File "bisect_m1_error.py", line 63, in <module>
    main()
  File "bisect_m1_error.py", line 59, in main
    reveal_error()
  File "bisect_m1_error.py", line 32, in reveal_error
    torch.onnx.export(conv, y, cpu_model_path, export_params=True, training=False)
  File "<@beartype(torch.onnx.utils.export) at 0x1281f5a60>", line 136, in export
  File "pytorch/venv/lib/python3.9/site-packages/beartype/_decor/_error/errormain.py", line 301, in raise_pep_call_exception
    raise exception_cls(  # type: ignore[misc]
beartype.roar.BeartypeCallHintParamViolation: @beartyped export() parameter training=False violates type hint <class 'torch._C._onnx.TrainingMode'>, as False not instance of <protocol "torch._C._onnx.TrainingMode">.
```

when `TORCH_ONNX_EXPERIMENTAL_RUNTIME_TYPE_CHECK` is not set and `beartype` is installed, a warning message is emitted.

```
>>> torch.onnx.export("foo", "bar", "f")
<stdin>:1: CallHintViolationWarning: Traceback (most recent call last):
  File "/home/justinchu/dev/pytorch/torch/onnx/_internal/_beartype.py", line 54, in _coerce_beartype_exceptions_to_warnings
    return beartyped(*args, **kwargs)
  File "<@beartype(torch.onnx.utils.export) at 0x7f1d4ab35280>", line 39, in export
  File "/home/justinchu/anaconda3/envs/pytorch/lib/python3.9/site-packages/beartype/_decor/_error/errormain.py", line 301, in raise_pep_call_exception
    raise exception_cls(  # type: ignore[misc]
beartype.roar.BeartypeCallHintParamViolation: @beartyped export() parameter model='foo' violates type hint typing.Union[torch.nn.modules.module.Module, torch.jit._script.ScriptModule, torch.jit.ScriptFunction], as 'foo' not <protocol "torch.jit.ScriptFunction">, <protocol "torch.nn.modules.module.Module">, or <protocol "torch.jit._script.ScriptModule">.

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/justinchu/dev/pytorch/torch/onnx/_internal/_beartype.py", line 63, in _coerce_beartype_exceptions_to_warnings
    return func(*args, **kwargs)
  File "/home/justinchu/dev/pytorch/torch/onnx/utils.py", line 482, in export
    _export(
  File "/home/justinchu/dev/pytorch/torch/onnx/utils.py", line 1422, in _export
    with exporter_context(model, training, verbose):
  File "/home/justinchu/anaconda3/envs/pytorch/lib/python3.9/contextlib.py", line 119, in __enter__
    return next(self.gen)
  File "/home/justinchu/dev/pytorch/torch/onnx/utils.py", line 177, in exporter_context
    with select_model_mode_for_export(
  File "/home/justinchu/anaconda3/envs/pytorch/lib/python3.9/contextlib.py", line 119, in __enter__
    return next(self.gen)
  File "/home/justinchu/dev/pytorch/torch/onnx/utils.py", line 95, in select_model_mode_for_export
    originally_training = model.training
AttributeError: 'str' object has no attribute 'training'
```

We see the error is caught right when the type mismatch happens, improving from what otherwise would become `AttributeError: 'str' object has no attribute 'training'`
Pull Request resolved: https://github.com/pytorch/pytorch/pull/83673
Approved by: https://github.com/BowenBao
2022-08-25 21:24:37 +00:00
BowenBao
daca0ee5e2 [ONNX] Introduce ONNXScopeName (#82038)
Update `_setup_trace_module_map` to always record module/layer info
in `Scope` attribute for nodes.
Extend `Scope` name to not only record module typename, but also
module object variable name. Both names are formatted and stored
as `name` attribute in `Scope`.
Introduce `ONNXScopeName` class to manage the formatting and parsing.
Updated local function export code adjusting to this update.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/82038
Approved by: https://github.com/AllenTiTaiWang, https://github.com/justinchuby, https://github.com/abock, https://github.com/malfet
2022-08-22 20:49:21 +00:00
Justin Chu
e4f74f0891 [ONNX] Update the default opset to version 14 (#83284)
Update the default opset by running the `update_default_opset_version.py` script. The update is done in a regularly to ensure we are in sync with the onnx updates. All changes are produced by the script.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/83284
Approved by: https://github.com/AllenTiTaiWang, https://github.com/malfet, https://github.com/BowenBao
2022-08-18 19:13:38 +00:00
BowenBao
017ecb782d [ONNX] Update legacy code, initialize onnx_shape_inference=True by default (#82767)
Legacy code has onnx_shape_inference=False by default, which is misleading
as every other export api sets it to True unless otherwise overriden by caller.
There is only two tests that need updating according to this change.
* test_utility_funs.py::test_constant_fold_shape. The resulting number of nodes
  in graph is increased by 1, due to that previously the extra constant node was
  added as initializer.
* test_utility_funs.py::test_onnx_function_substitution_pass. Enabling onnx
  shape inference discovered discrepancy in test input shape and supplied dynamic
  axes arguments.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/82767
Approved by: https://github.com/justinchuby, https://github.com/abock
2022-08-10 21:50:13 +00:00
Justin Chu
f5701a1f9a [ONNX] Remove unused patching methods (#83006)
### Description
<!-- What did you change and why was it needed? -->

Remove unused patching methods:

- `torch._C.Graph.constant`
- unpatch `torch._C.Node.__getitem__` and move the helper function to `symbolic_helper`

Add typing annotations

### Issue
<!-- Link to Issue ticket or RFP -->

#76254

### Testing
<!-- How did you test your change? -->

Unit tested
Pull Request resolved: https://github.com/pytorch/pytorch/pull/83006
Approved by: https://github.com/BowenBao
2022-08-09 19:24:03 +00:00
qqaatw
9b4dc56c83 [ONNX] Fix quantization outputs' dtype (#79690)
Part of #79263

Previously, all quantized PyTorch tensors are all casted to the dtypes which comply with ONNX's definition, i.e. `scale` is casted to `double`, and `zero_point` is casted to `int64`. These casts lead to inconsistent dtypes when comparing PyTorch's outputs and ONNX runtime's outputs.

Now, `cast_onnx_accepted` argument is added to `unpack_quantized_tensor` function. When making example inputs for ONNX, we cast them to the ONNX compliant dtypes; otherwise, they are casted to PyTorch default types for quantization.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/79690
Approved by: https://github.com/justinchuby, https://github.com/BowenBao
2022-08-09 18:32:03 +00:00
shubhambhokare1
95d873855e [ONNX] Inline prim::PythonOp for Autograd Function Export (#74765)
Add flag (inline_autograd) to enable inline export of model consisting of autograd functions. Currently, this flag should only be used in TrainingMode.EVAL and not for training.

An example:

If a model containing ``autograd.Function`` is as follows
```
                class AutogradFunc(torch.autograd.Function):
                  @staticmethod
                  def forward(ctx, i):
                      result = i.exp()
                      result = result.log()
                      ctx.save_for_backward(result)
                      return result
```
Then the model is exported as
```
                graph(%0 : Float):
                  %1 : Float = ^AutogradFunc(%0)
                  return (%1)
```
If inline_autograd is set to True, this will be exported as
```
                graph(%0 : Float):
                  %1 : Float = onnx::Exp(%0)
                  %2 : Float = onnx::Log(%1)
                  return (%2)
```

If one of the ops within the autograd module is not supported, that particular node is exported as is mirroring ONNX_FALLTHROUGH mode

Fixes: #61813
Pull Request resolved: https://github.com/pytorch/pytorch/pull/74765
Approved by: https://github.com/BowenBao, https://github.com/malfet
2022-08-03 23:30:19 +00:00
Huy Do
6ea422dd0b Format torch/onnx with ufmt (#82137)
This is the last batch for the new ufmt (black + usort) linter. After this, black linter can finally be replaced. The previous PR to format ONNX tests was #81335
Pull Request resolved: https://github.com/pytorch/pytorch/pull/82137
Approved by: https://github.com/kit1980, https://github.com/AllenTiTaiWang
2022-07-25 22:42:21 +00:00
Justin Chu
d1d2687d34 [ONNX] Fix potentially unbound variables (#79789)
Pylint alerts that some variables may be unbound. This PR fixes the errors.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/79789
Approved by: https://github.com/garymm
2022-06-29 17:01:49 +00:00
Justin Chu
4b817f5816 [ONNX] Improve docstrings and typing annotations (#78231)
- Remove wrappers in `__init__` around utils and instead expose those functions directly. Move the docstrings from `__init__` to corresponding functions in utils
- Annotate `torch.onnx.export` types
- Improve docstrings
Pull Request resolved: https://github.com/pytorch/pytorch/pull/78231
Approved by: https://github.com/BowenBao
2022-06-16 02:59:24 +00:00
Justin Chu
c8b9b6266b [ONNX] Fix arg type in _set_training_mode (#78583)
When `TrainingMode.PRESERVE` is set for export, the exporter used to change the model's training mode based on some logic. Now we respect the option and not touch the model's training state.

- Previously `_set_training_mode`'s behavior doesn't match what the global variable expects. This PR removes the deprecated `_set_training_mode` and makes the type correct.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/78583
Approved by: https://github.com/BowenBao
2022-06-15 23:47:12 +00:00
Justin Chu
d3ef5c3fa3 [ONNX] Clean up __init__ in torch.onnx (#78446)
- Move definitions in `__init__` to internal classes and expose them by importing to init (prevent circular dependencies): https://github.com/pytorch/pytorch/wiki/torch.onnx-Namespacing
  - Context classes and enums are moved to `_exporter_states.py`
  - Exceptions are moved to `errors.py`
- Define `__all__` for torch.onnx. https://github.com/pytorch/pytorch/wiki/Public-API-definition-and-documentation
- Moved `utils.__IN_ONNX_EXPORT` to `GLOBALS.in_onnx_export`
- Deprecated `torch.onnx._export`

Precedes #78231

Using this as an aid for finding public functions:

```python
list(filter(lambda x: not x.startswith("_"), torch.onnx.utils.__dict__.keys()))
```
Pull Request resolved: https://github.com/pytorch/pytorch/pull/78446
Approved by: https://github.com/BowenBao
2022-06-14 04:35:06 +00:00
BowenBao
530dcc2b94 [ONNX] Tool to verify exported model discrepancy between sets of inputs
A graph is exported for each set of inputs. The exported graphs are then compared
to each other, and discrepancies are reported. This function first checks the jit
graph, and then the onnx graph.

Unless otherwise specified, the jit/ONNX graph is expected to be the same, regardless
of the inputs it used for exporting. A discrepancy would imply the graph exported is
not accurate when running with other set of inputs, which will typically results in
runtime error or output mismatches.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/78323

Approved by: https://github.com/justinchuby, https://github.com/garymm
2022-06-06 20:29:20 +00:00
Justin Chu
161e931156 [ONNX] Modernize python syntax (#77935)
Use pyupgrade(https://github.com/asottile/pyupgrade) and flynt to modernize python syntax

```sh
pyupgrade --py36-plus --keep-runtime-typing torch/onnx/**/*.py
pyupgrade --py36-plus --keep-runtime-typing test/onnx/**/*.py
flynt torch/onnx/ --line-length 120
```

- Use f-strings for string formatting
- Use the new `super()` syntax for class initialization
- Use dictionary / set comprehension
Pull Request resolved: https://github.com/pytorch/pytorch/pull/77935
Approved by: https://github.com/BowenBao
2022-05-24 22:52:37 +00:00
Justin Chu
0d76299ff7 [ONNX] Clean up module imports (#77423)
Cleaning up onnx module imports to prepare for updating `__init__`.

- Simplify importing the `_C` and `_C._onnx` name spaces
- Remove alias of the symbolic_helper module in imports
- Remove any module level function imports. Import modules instead
    - Alias `symbilic_opsetx` as `opsetx`
- Fix some docstrings

Requires:
- https://github.com/pytorch/pytorch/pull/77448
Pull Request resolved: https://github.com/pytorch/pytorch/pull/77423
Approved by: https://github.com/BowenBao
2022-05-20 01:56:24 +00:00
Justin Chu
563c2719bf [ONNX] Refactor to remove inline imports - attempt 2 (#77448)
Re-land
- #77142

(diff: https://github.com/pytorch/pytorch/compare/c08b8f0..justinchuby:justinchu/remove-patch2)

Fixed:
- Delay import symbolic_opsets in the registry.

Tested locally with torchvision
Pull Request resolved: https://github.com/pytorch/pytorch/pull/77448
Approved by: https://github.com/garymm
2022-05-16 14:44:24 +00:00
PyTorch MergeBot
6b366dd3c1 Revert "[ONNX] Refactor to remove inline imports (#77142)"
This reverts commit c08b8f0967.

Reverted https://github.com/pytorch/pytorch/pull/77142 on behalf of https://github.com/malfet
2022-05-13 19:44:17 +00:00
Justin Chu
c08b8f0967 [ONNX] Refactor to remove inline imports (#77142)
Reduce circular dependencies

- Lift constants and flags from `symbolic_helper` to `_constants` and `_globals`
    - Standardized constant naming to make it consistant
- Make `utils` strictly dependent on `symbolic_helper`, removing inline imports from symbolic_helper
- Move side effects from `utils` to `_patch_torch`

Pull Request resolved: https://github.com/pytorch/pytorch/pull/77142
Approved by: https://github.com/garymm, https://github.com/BowenBao
2022-05-13 03:46:33 +00:00
Justin Chu
78d3798181 [ONNX] Fix type comparison in utils._need_symbolic_context (#77365)
In `_need_symbolic_context`, when the annotation is postponed evaluated, the annotation is a string and not a type. We need to use get_type_hints to get the real type.

For example,

```python
def g(a: int) -> int: return a

def f(a: "int") -> "int": return a
```

we will get the correct type `int` for both g and f with `typing.get_type_hints`. Otherwise, the type for `a` in `f` will be a string and is not comparable to the type `int` - `issubclass` will complain.

This is necessary as we will use postponed typing evaluation to break circular dependencies.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/77365
Approved by: https://github.com/BowenBao
2022-05-12 23:49:04 +00:00
BowenBao
93953a48b7 [ONNX] Bug fix: opset_version checked before set (#76928)
Pull Request resolved: https://github.com/pytorch/pytorch/pull/76928
Approved by: https://github.com/justinchuby, https://github.com/garymm
2022-05-11 17:16:22 +00:00
Justin Chu
5dd1c67776 [ONNX] Format ONNX python with black
Format all onnx python code with black and isort with

```sh
isort torch/onnx/ test/onnx
black torch/onnx/ test/onnx
```

Updated lintrunner config to include these paths.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/76754
Approved by: https://github.com/suo, https://github.com/BowenBao
2022-05-05 00:19:22 +00:00
BowenBao
679fc90cdb [ONNX] Support optional type (#68793) (#73284)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/73284

Some important ops won't support optional type until opset 16,
so we can't fully test things end-to-end, but I believe this should
be all that's needed. Once ONNX Runtime supports opset 16,
we can do more testing and fix any remaining bugs.

Test Plan: Imported from OSS

Reviewed By: albanD

Differential Revision: D34625646

Pulled By: malfet

fbshipit-source-id: 537fcbc1e9d87686cc61f5bd66a997e99cec287b

Co-authored-by: BowenBao <bowbao@microsoft.com>
Co-authored-by: neginraoof <neginmr@utexas.edu>
Co-authored-by: Nikita Shulga <nshulga@fb.com>
(cherry picked from commit 822e79f31ae54d73407f34f166b654f4ba115ea5)
2022-05-04 20:24:30 +00:00
Justin Chu
2e841e68b2 [ONNX] Documentation and typing annotations in registry
Updating the docstrings and type annotations as I walk through the code.

- Turned some comments into docstrings.
- Added type annotations for some functions in utils and the registry
- Removed direct function imports; importing functions makes name space collision easier to happen and refactoring/code analysis harder: https://google.github.io/styleguide/pyguide.html#22-imports
- Formatted touched files with black
Pull Request resolved: https://github.com/pytorch/pytorch/pull/76255
Approved by: https://github.com/BowenBao
2022-04-28 18:24:24 +00:00
Thiago Crepaldi
90d31cb311 Emit ATen ops when symbolics raise + minor fixes
Currently `torch.onnx.export(.., operator_export_type=OperatorExportTypes.ONNX_ATEN_FALLBACK)` only issues ATen ops through explicit requests (e.g. `g.at()`) calls inside each op symbolic function. This is done based on specific conditions such as `operator_export_type==OperatorExportTypes.ONNX_ATEN_FALLBACK)` or `is_caffe2_aten_fallback()`

This PR extends the ATen fallback mechanism for scenarios when the symbolic function raises `RuntimeError` during export. The idea is that partial implementation of existing ONNX ops can fallback to ATen as a last resort. That is valuable because each operator can have many input combinations and not all are always implemented.

A minor fix was done to make sure the `overload_name` attribute is added to explicit ATen op fallback requests when a symbolic is not registered to a particular op.

ps: The behavior for builds with BUILD_CAFFE2=1 is not changed to ensure BC.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/74759
Approved by: https://github.com/garymm, https://github.com/msaroufim
2022-04-23 21:24:25 +00:00
BowenBao
2c748b7573 [ONNX] Trace model if quantization is detected
Previously pre-tracing model is required for exporting quantized model.
e.g. calling `traced_m = torch.jit.trace(model, inputs)` and export `traced_m`.
The reason was quantized weights are stored in a unique `PackedParam` structure,
and they need to be handled by tracing to be exportable.
This PR enables export api to call tracing underneath if it detects quantization
in the model.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/75921

Approved by: https://github.com/garymm
2022-04-22 17:27:32 +00:00
Thiago Crepaldi
9bbe1d632e Fix ONNX ATen fallback for non-caffe2 engines
This PR introduces 3 BC changes:

First, this PR propagates `BUILD_CAFFE2` flag to `libtorch` and `libtorch_python`, which is necessary for non-caffe2 ONNX runtimes when using `ONNX_ATEN_FALLBACK` operator export type.

Second, as a complement of https://github.com/pytorch/pytorch/pull/68490, this PR refactors Caffe2's Aten ops symbolics to consider not only the `operator_export_type` (aka `ONNX_ATEN_FALLBACK`) to emit Caffe2 Aten ops, but also whether `BUILD_CAFFE2` (which is called `torch.onnx._CAFFE2_ATEN_FALLBACK` in python binding) is set.

Lastly, it renames `onnx::ATen` to `aten::ATen` for ONNX spec consistency in a BC fashion.
ONNX doesn't have `ATen` op on its spec, but PyTorch ONNX converter emits them. Non-Caffe2 backend engines would be mislead by such operator's name/domain. A non-ideal workaround would be to have Aten ops handled based on its name and ignore the (non-complaint) domain. Moreover, users could incorrectly file bugs to either ONNX or ONNX Runtime when they inspect the model and notice the presence of an unspecified ONNX operator.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/73954
Approved by: https://github.com/BowenBao, https://github.com/malfet, https://github.com/garymm, https://github.com/jiafatom
2022-04-14 23:18:45 +00:00
BowenBao
144b7de9dd [ONNX] Adjust is_train flag for onnx pass deduplicate initializers
Previous logic didn't consider the case for TrainingMode.PRESERVE.
A more direct way is to check `model.training`, which is the accurate
training mode, set by `exporter_context(model, training)`.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/74247
Approved by: https://github.com/garymm
2022-03-22 22:39:20 +00:00
BowenBao
54a6942f8d [ONNX] ONNX Exporter logging (#71342)
Summary:
Add ONNX exporter logging facility. Supporting both C++/Python logging api. Logging can be turned on/off. Logging output stream can be either set to `stdout` or `stderr`.

A few other changes:
* When exception is raised in passes, the current IR graph being processed will be logged.
* When exception is raised from `_jit_pass_onnx` (the pass that converts nodes from namespace `ATen` to `ONNX`), both ATen IR graph and ONNX IR graph under construction will be logged.
* Exception message for ConstantFolding is truncated to avoid being too verbose.
* Update the final printed IR graph with node name in ONNX ModelProto as node attribute. Torch IR Node does not have name. Adding this to printed IR graph helps debugging.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/71342

Reviewed By: msaroufim

Differential Revision: D34433473

Pulled By: malfet

fbshipit-source-id: 4b137dfd6a33eb681a5f2612f19aadf5dfe3d84a
(cherry picked from commit 67a8ebed5192c266f604bdcca931df6fe589699f)
2022-03-17 19:40:03 +00:00
BowenBao
9210e8f540 [ONNX] Adds overload_name to Aten op (#69378) (#73280)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/73280

This PR adds a new attribute overload_name to the Aten node so that third party applications can implement calls to libtorch without using PyTorch source code.

This is necessary because torch's torch::jit::findOperatorFor(fullname) requires a full name, including operator and overload names.

ATen op was originally created for Caffe2, which leveraged the availability of the pytorch yaml files to create calls to the aten oeprators directly, not relying on torch::jit::findOperatorFor

The first part of the PR refactors all symbolics that create Aten ops, so that there is a single helper for this operator. Next all symbolics are updated to pass in the relevant overload name, if empty string is not applicable

Test Plan: Imported from OSS

Reviewed By: jbschlosser

Differential Revision: D34625645

Pulled By: malfet

fbshipit-source-id: 37d58cfb5231833768172c122efc42edf7d8609a
(cherry picked from commit e92f09117d3645b38bc3235b30aba4b4c7c71dfa)
2022-03-09 14:26:18 +00:00
BowenBao
b3cfc74f0f [ONNX] Capture annotated attributes for local function
Enables local function export to capture annotated attributes.
For example:
```python
class M(torch.nn.Module):
    num_layers: int

    def __init__(self, num_layers):
        super().__init__()
        self.num_layers = num_layers

    def forward(self, args):
        ...
```
`num_layers` will now be captured as attribute of local function `M`.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/72883
2022-02-28 18:56:18 +00:00
BowenBao
28bf2f80cf Don't call _jit_pass_onnx_function_extraction if export_modules_as_functions is False (#69742)
* fix clang-format violations

* Don't call _jit_pass_onnx_function_extraction if export_modules_as_functions is False

It's just wasteful.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/73100
2022-02-22 22:43:53 +00:00
BowenBao
2791725a84 Integrate full ONNX check into ONNX export API (#71125)
Pull Request resolved: https://github.com/pytorch/pytorch/pull/72988
2022-02-18 18:40:09 +00:00
BowenBao
32f6a1e2a2 [ONNX] First version of quantized model export: Support quantized.Linear (#69232)
Co-authored-by: David Fan <jiafamicrosoft.com>

Pull Request resolved: https://github.com/pytorch/pytorch/pull/72986
2022-02-18 18:27:26 +00:00
BowenBao
cc792746d2 [ONNX] De-duplicate initializers (#68202) (#69547)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/69547

ScriptModule export introduces duplicated ONNX initializers for shared weights, unnecessarily increases ONNX model size. This PR de-duplicates ONNX initializers for model exported in eval mode, by checking if the underlying tensors share the same `data_ptr`, `strides` and `sizes`.

Test Plan: Imported from OSS

Reviewed By: msaroufim

Differential Revision: D32994271

Pulled By: malfet

fbshipit-source-id: 10ac66638b6255890875272472aa9ed07a5b1d9a

Co-authored-by: BowenBao <bowbao@microsoft.com>
(cherry picked from commit d7cbde940c)
2022-02-11 22:05:15 +00:00
BowenBao
04c5d978b9 [ONNX] Refactor _run_symbolic_function (#67573) (#68491)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/68491

* Allows implementing symbolic functions for domains other than `aten`, for example `prim`, in symbolic_opset#.py.
* Allows symbolic function to access extra context if needed, through `SymbolicFunctionState`.
  * Particularly, the `prim::PythonOp` special case can access node without the need of passing node through inputs. Updates will be made downstreams, and in a follow-up PR we will remove the previous workaround in exporter.
* `prim::Loop`, `prim::If`, etc are now moved outside of `_run_symbolic_function` from utils.py, and to symbolic_opset9.py.

Motivation for this change:
- Better maintainability and reducing complexity. Easier to add symbolic for operators, both simple and complex ones (that need additional context), without the former needing to know the existence of the latter.
- The design idea was long outdated. prim ops are no longer rare special cases, and they shouldn't all be handled inside `_run_symbolic_function`. As a result this function becomes too clumsy. There were also prim ops symbolic added in symbolic_opset#.py with signature `prim_[opname]`, creating separation and confusion.

Test Plan: Imported from OSS

Reviewed By: jansel

Differential Revision: D32483782

Pulled By: malfet

fbshipit-source-id: f9affc31b1570af30ffa6668da9375da111fd54a

Co-authored-by: BowenBao <bowbao@microsoft.com>
(cherry picked from commit 1e04ffd2fd)
2022-02-11 18:35:35 +00:00
BowenBao
eb4238fc26 Allow caffe2-specific graph transformations for OperatorExportTypes.ONNX_ATEN_FALLBACK when BUILD_CAFFE2 is ON (#67460) (#68490)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/68490

The use of ATEN as a fallback operator during ONNX conversion is important for increasing operator coverage or even provide more efficient implementations over some ONNX ops.

Currently this feature is available through `OperatorExportTypes.ONNX_ATEN_FALLBACK`,
but it also performs changes to the graph that are runnable by Caffe2, only.

This PR introduces restricts caffe2-specific graph transformations for `ONNX_ATEN_FALLBACK`
operator export type for when pytorch is built with caffe2 support (aka BUILD_CAFFE2=1 during build)

The first version of this PR introduced a new operator export type `ONNX_ATEN__STRICT_FALLBACK`,
which essentially is the same as `ONNX_ATEN_FALLBACK` but without caffe2 transformations.
It was preferred to not introduce a new operator export type, but to refine the existing aten fallback one

## BC-breaking note
### The global constant `torch.onnx.PYTORCH_ONNX_CAFFE2_BUNDLE` is removed in favor of
a less visible `torch.onnx._CAFFE2_ATEN_FALLBACK`.
`PYTORCH_ONNX_CAFFE2_BUNDLE` is really a dead code flag always set to False.
One alternative would be fixing it, but #66658 disables Caffe2 build by default.
Making a Caffe2 feature a private one seems to make more sense for future deprecation.

### The method `torch.onnx.export` now defaults to ONNX when `operator_export_type` is not specified.
Previously `torch.onnx.export's operator_export_type` intended to default to `ONNX_ATEN_FALLBACK` when `PYTORCH_ONNX_CAFFE2_BUNDLE` was set, but it would never happen as `PYTORCH_ONNX_CAFFE2_BUNDLE` is always undefined

 Co-authored-by: Nikita Shulga <nshulga@fb.com>

Test Plan: Imported from OSS

Reviewed By: jansel

Differential Revision: D32483781

Pulled By: malfet

fbshipit-source-id: e9b447db9466b369e77d747188685495aec3f124
(cherry picked from commit 5fb1eb1b19)
2022-02-10 03:26:48 +00:00
BowenBao
cf70466970 [ONNX] Improve scope inference in function extraction
Cover more cases of scope inferencing where consecutive nodes don't have valid scope information. Usually these nodes are created in some pass where authors forgot to assign meaningful scope to them.
* One rule of `InferScope` is to check if the current node's outputs' users share the same scope. Recursively run `InferScope` on the user nodes if they are missing scope as well. Since the graph is SSA, the depth is finite.
* Fix one pass that missed scope information for a new node.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/71897
2022-01-31 23:58:53 +00:00
BowenBao
804f13289e [ONNX] Update opset_version restriction for local function
Export should fail if export_modules_as_functions is set and opset_version<15.
This is because opeset_version < 15 implies IR version < 8, which means no local function support.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/71619
2022-01-27 00:21:13 +00:00
Emilio Castillo
e2dc2aca93 Export ONNX models with readable input/output names (#68976)
Summary:
For some ONNX exported models, the inputs/outputs names have sometimes a numeric value and this makes pretty hard to inspect the generated graphs in the case of large models.

The solution in this PR was initially submitted to our internal utilities library by take-cheeze https://github.com/pfnet/pytorch-pfn-extras/pull/102

Now we would like to upstream this change by adding an extra kwarg when exporting the model to allow replacing these numeric names with actual debuggable ones.

As an example, the following code shows that the module output is `3`

```python
g, p, o = _model_to_graph(module, torch.ones(1, 10))
for n in g.nodes():
    for v in n.outputs():
        print(v.debugName())
```
output
```
3
```

With this PR

```
v3_Gemm
```

This allows identifying this out as a value from the associated Gemm layer.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/68976

Reviewed By: jansel

Differential Revision: D33662246

Pulled By: msaroufim

fbshipit-source-id: 45f56eef2a84d9a318db20c6a6de6c2743b9cd99
(cherry picked from commit 513c1d28f1)
2022-01-21 00:34:56 +00:00
BowenBao
ff78c73286 [ONNX] Remove f arg from export_to_pretty_string (#69045) (#69546)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/69546

The arg is not used and was previously deprecated.

Also remove torch.onnx._export_to_pretty_string. It's redundant with the
public version.

Test Plan: Imported from OSS

Reviewed By: malfet

Differential Revision: D32994270

Pulled By: msaroufim

fbshipit-source-id: f8f3933b371a0d868d9247510bcd73c31a9d6fcc
2022-01-12 21:31:36 -08:00
Deyu Huang
d32efe8bc2 [ONNX] Remove the argument use_external_data_format of export() method entirely. (#67080) (#67811)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/67811

* remove the argument use_external_data_format of export() method entirely

Test Plan: Imported from OSS

Reviewed By: msaroufim

Differential Revision: D32181302

Pulled By: malfet

fbshipit-source-id: 4bc1448b7487bb9dfdad4e36008ff5b227fd64a3

Co-authored-by: hwangdeyu <dejack953@outlook.com>
2021-11-15 17:20:04 -08:00
Thiago Crepaldi
9d25554d45 [ONNX] Allow registration of custom symbolics for aten namespace (#66481) (#67810)
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/67810

Test Plan: Imported from OSS

Reviewed By: msaroufim

Differential Revision: D32181303

Pulled By: malfet

fbshipit-source-id: af2a715dc554b958fa3f5a7a8ae96cb3f7d112bb
2021-11-15 17:18:39 -08:00
Deyu Huang
48c8de45b0 [ONNX] Remove the argument example_outpus of export() method entirely. (#67082) (#67809)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/67809

* remove the argument example_outpus of export() method entirely

[ONNX] Follow-up: Remove the argument example_outpus of export() method entirely. (#67629)

* Resolve CI failure

* remove test after removing example_outputs

[ONNX] Follow-up: Follow-up: Remove the argument example_outpus of export() method entirely (#67719)

Removing unused import, resolving flake error.

Test Plan: Imported from OSS

Reviewed By: msaroufim

Differential Revision: D32181305

Pulled By: malfet

fbshipit-source-id: ba00547b7cb455ace86606b1bda643c02bdcfa1b

Co-authored-by: hwangdeyu <dejack953@outlook.com>
2021-11-12 17:06:26 -08:00
Bowen Bao
02e35ce17b [ONNX] Update onnx function export with comments and clean up (#66817) (#67803)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/67803

* Addresses comments from #63589

[ONNX] remove torch::onnx::PRODUCER_VERSION (#67107)

Use constants from version.h instead.
This simplifies things since we no longer have to update
PRODUCER_VERSION for each release.

Also add TORCH_VERSION to version.h so that a string is available for
this purpose.

[ONNX] Set `ir_version` based on opset_version. (#67128)

This increases the odds that the exported ONNX model will be usable.
Before this change, we were setting the IR version to a value which may
be higher than what the model consumer supports.

Also some minor clean-up in the test code:
* Fix string replacement.
* Use a temporary file so as to not leave files around in the test
  current working directory.

Test Plan: Imported from OSS

Reviewed By: msaroufim

Differential Revision: D32181306

Pulled By: malfet

fbshipit-source-id: 02f136d34ef8f664ade0bc1985a584f0e8c2b663

Co-authored-by: BowenBao <bowbao@microsoft.com>
Co-authored-by: Gary Miguel <garymiguel@microsoft.com>
Co-authored-by: Nikita Shulga <nshulga@fb.com>
2021-11-05 10:35:35 -07:00
Jay Zhang
26241994b2 Remove the argument strip_doc_string of export() method entirely. (#66615) (#67278)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/67278

Remove the argument strip_doc_string of export() method entirely.

Test Plan: Imported from OSS

Reviewed By: msaroufim

Differential Revision: D31962512

Pulled By: malfet

fbshipit-source-id: 168ad3f157a80d1edd7a9053783b3f3deb2ecf43

Co-authored-by: fatcat-z <jiz@microsoft.com>
2021-10-28 19:25:07 -07:00
Jay Zhang
43d51254bf Deprecate the argument _retain_param_name of export() method entirely. (#66617) (#67277)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/67277

Remove the argument _retain_param_name of export() method entirely.

Test Plan: Imported from OSS

Reviewed By: msaroufim

Differential Revision: D31962514

Pulled By: malfet

fbshipit-source-id: 8ac5e3a4a7821cc580951a7f167fd20069116350

Co-authored-by: fatcat-z <jiz@microsoft.com>
2021-10-28 19:25:05 -07:00
Jay Zhang
40920185ac [ONNX] Remove the argument enable_onnx_checker of export() method entirely. (#66611) (#67276)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/67276

[ONNX] Remove argument _retain_param_name from torch.onnx.export() function.

Test Plan: Imported from OSS

Reviewed By: msaroufim

Differential Revision: D31962520

Pulled By: malfet

fbshipit-source-id: 86ee15f525261c0da74175e47dd74eeb169ac47f

Co-authored-by: fatcat-z <jiz@microsoft.com>
2021-10-28 19:25:03 -07:00
Nikita Shulga
b18c298f24 ONNX: Delete or document skipped ORT tests (#64470) (#66143)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/66143

Delete test_list_remove. There's no point in testing conversion of
this model since TorchScript doesn't support it.

Add a link to an issue tracking test_embedding_bag_dynamic_input.

[ONNX] fix docs (#65379)

Mainly fix the sphinx build by inserting empty before
bulleted lists.

Also some minor improvements:
Remove superfluous descriptions of deprecated and ignored args.
The user doesn't need to know anything other than that they are
deprecated and ignored.

Fix custom_opsets description.

Make indentation of Raises section consistent with Args section.

[ONNX] publicize func for discovering unconvertible ops (#65285)

* [ONNX] Provide public function to discover all unconvertible ATen ops

This can be more productive than finding and fixing a single issue at a
time.

* [ONNX] Reorganize test_utility_funs

Move common functionality into a base class that doesn't define any
tests.

Add a new test for opset-independent tests. This lets us avoid running
the tests repeatedly for each opset.

Use simple inheritance rather than the `type()` built-in. It's more
readable.

* [ONNX] Use TestCase assertions rather than `assert`

This provides better error messages.

* [ONNX] Use double quotes consistently.

[ONNX] Fix code block formatting in doc (#65421)

Test Plan: Imported from OSS

Reviewed By: jansel

Differential Revision: D31424093

fbshipit-source-id: 4ced841cc546db8548dede60b54b07df9bb4e36e
2021-10-22 13:46:16 -07:00
Nikita Shulga
53a163a015 [ONNX] Export nn.Module call as ONNX local function (#63589) (#66140)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/66140

* Add new argument to export api to enable users specifying `nn.Module` classes that they wish to be exported as local function in ONNX model.
* Refactor `torch/csrc/jit/serialization/export.cpp`, and remove redundant `EncoderBase` class.
* ~~Contains changes from #63268~~
* Depends on #63716 to update onnx submodule.

Test Plan: Imported from OSS

Reviewed By: jansel

Differential Revision: D31424098

fbshipit-source-id: c949d0b01c206c30b4182c2dd1a5b90e32b7a0d3

Co-authored-by: BowenBao <bowbao@microsoft.com>
2021-10-22 13:44:56 -07:00
BowenBao
1cf317b85f [ONNX] Support exporting with Apex O2 (#65374) (#66700)
Summary:
Apex O2 hook state_dict to return fp16 weights as fp32. Exporter cannot identify them as same tensors.
Since this hook is only used by optimizer, it is safe to remove this hook while exporting.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/66700

Reviewed By: zou3519

Differential Revision: D31695132

Pulled By: malfet

fbshipit-source-id: 977bdf57240002498f3ad0f1a8046c352e9860e6
2021-10-18 11:54:09 -07:00
Gary Miguel
2506baf9c2 [ONNX] move CheckerError from torch.onnx.utils to torch.onnx (#66644)
Summary:
This moves it to where the user would expect it to be based on the
documentation and all the other public classes in the torch.onnx module.

Also rename it from ONNXCheckerError, since the qualified name
torch.onnx.ONNXCheckerError is otherwise redundant.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/66644

Reviewed By: malfet

Differential Revision: D31662559

Pulled By: msaroufim

fbshipit-source-id: bc8a57b99c2980490ede3974279d1124228a7406
2021-10-15 10:38:56 -07:00
Edward Yang
11bc435622 Allow registration of custom symbolics for prim namespace (#64460) (#66139)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/66139

[ONNX] Add prim::PythonOp check back in export.cpp (#64944)

Add prim::PythonOp check back in export.cpp

Test Plan: Imported from OSS

Reviewed By: malfet

Differential Revision: D31424102

fbshipit-source-id: 6d2eef767fab846ed79ea509e97b714072bac9f4

Co-authored-by: jiafatom <jiafa@microsoft.com>
2021-10-08 07:41:06 -07:00
Edward Yang
53fefaa916 [ONNX] Fix duplicated output same name case (#64190) (#66137)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/66137

* fix duplicated output node same output name issue.

Test Plan: Imported from OSS

Reviewed By: malfet

Differential Revision: D31424100

fbshipit-source-id: b1b06a92c51744030788b651f3a597d987a8deda

Co-authored-by: hwangdeyu <dejack953@outlook.com>
2021-10-08 07:41:01 -07:00
BowenBao
2d61009f4a [ONNX] Fix input sequence for pad op (#60554) (#64377)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/64377

* Fix for input primitive sequence

* Test mypy

* Fix for tracing tuples

* Fix for extra inputs

* flake8

* Rebase

* Fix for tracing tuples

Test Plan: Imported from OSS

Reviewed By: jansel

Differential Revision: D30919606

Pulled By: malfet

fbshipit-source-id: a718c4a12cda77b968cb636acd7aa63d7b5ba326
2021-09-30 21:08:45 -07:00
BowenBao
20143bf07f [ONNX] Deprecate use_external_data_format param from torch.onnx.export() function. (#62257) (#64382)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/64382

* This `use_external_data_format` parameter is used for large models cannot be exported because of the 2GB protobuf limit.

* When `use_external_data_format` set to True, the model is exported in ONNX external data format, in which case some of the model parameters are stored in external binary files and not in the ONNX model file itself.

* This PR will set this paramter to DEPRECATED and check the model proto sizes by code instead of by user, if the sizes lager than 2GB, then `use_external_data_format = True` automatically.

Test Plan: Imported from OSS

Reviewed By: ezyang

Differential Revision: D30905265

Pulled By: malfet

fbshipit-source-id: 82b4e17bfa6a8de2bfd700a5282c12f6835603cb

Co-authored-by: hwangdeyu <dejack953@outlook.com>
2021-09-23 22:20:48 -07:00
BowenBao
478d4cf883 [ONNX] Deprecated the example_outputs param from torch.onnx.export() function. (#62815) (#64380)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/64380

* `example_outputs` used to determine the type and shape of the outputs without tracing the execution of the model. And it must be provided when exporting a ScriptModule or ScriptFunction when using export() function.

* Since we can work out `example_outputs` in internal function instead of being provided by user, so we deprecated this argument in the export() function to increase user experience of calling this function.

Test Plan: Imported from OSS

Reviewed By: ezyang

Differential Revision: D30905266

Pulled By: malfet

fbshipit-source-id: d00b00d7d02b365d165028288ad915678caa51f2

Co-authored-by: hwangdeyu <dejack953@outlook.com>
2021-09-23 22:20:46 -07:00