Skip to content

Commit 69df5c6

Browse files
committed
update api
1 parent fbd2583 commit 69df5c6

File tree

2 files changed

+20
-9
lines changed

2 files changed

+20
-9
lines changed

beginner_source/onnx/export_simple_model_to_onnx_tutorial.py

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -7,12 +7,12 @@
77
Export a PyTorch model to ONNX
88
==============================
99
10-
**Author**: `Ti-Tai Wang <https://github.com/titaiwangms>`_ and `Xavier Dupré <https://github.com/xadupre>`_
10+
**Author**: `Ti-Tai Wang <https://github.com/titaiwangms>`_ and Thiago Crepaldi <https://github.com/thiagocrepaldi>`_.
1111
1212
.. note::
13-
As of PyTorch 2.1, there are two versions of ONNX Exporter.
13+
As of PyTorch 2.5, there are two versions of ONNX Exporter.
1414
15-
* ``torch.onnx.dynamo_export`` is the newest (still in beta) exporter based on the TorchDynamo technology released with PyTorch 2.0
15+
* ``torch.onnx.export(..., dynamo=True)`` is the newest (still in beta) exporter based on the TorchDynamo technology released with PyTorch 2.0
1616
* ``torch.onnx.export`` is based on TorchScript backend and has been available since PyTorch 1.2.0
1717
1818
"""
@@ -21,7 +21,7 @@
2121
# In the `60 Minute Blitz <https://pytorch.org/tutorials/beginner/deep_learning_60min_blitz.html>`_,
2222
# we had the opportunity to learn about PyTorch at a high level and train a small neural network to classify images.
2323
# In this tutorial, we are going to expand this to describe how to convert a model defined in PyTorch into the
24-
# ONNX format using TorchDynamo and the ``torch.onnx.dynamo_export`` ONNX exporter.
24+
# ONNX format using TorchDynamo and the ``torch.onnx.export(..., dynamo=True)`` ONNX exporter.
2525
#
2626
# While PyTorch is great for iterating on the development of models, the model can be deployed to production
2727
# using different formats, including `ONNX <https://onnx.ai/>`_ (Open Neural Network Exchange)!
@@ -90,7 +90,16 @@ def forward(self, x):
9090

9191
torch_model = MyModel()
9292
torch_input = torch.randn(1, 1, 32, 32)
93-
onnx_program = torch.onnx.dynamo_export(torch_model, torch_input)
93+
onnx_program = torch.onnx.export(torch_model, torch_input, dynamo=True)
94+
95+
######################################################################
96+
# 3.5. (Optional) Optimize the ONNX model
97+
# ---------------------------------------
98+
#
99+
# The ONNX model can be optimized with constant folding, and elimination of redundant nodes.
100+
# The optimization is done in-place, so the original ONNX model is modified.
101+
102+
onnx_program.optimize()
94103

95104
######################################################################
96105
# As we can see, we didn't need any code change to the model.

beginner_source/onnx/intro_onnx.py

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
====================
88
99
Authors:
10-
`Ti-Tai Wang <https://github.com/titaiwangms>`_ and `Xavier Dupré <https://github.com/xadupre>`_
10+
`Ti-Tai Wang <https://github.com/titaiwangms>`_ and Thiago Crepaldi <https://github.com/thiagocrepaldi>`_.
1111
1212
`Open Neural Network eXchange (ONNX) <https://onnx.ai/>`_ is an open standard
1313
format for representing machine learning models. The ``torch.onnx`` module provides APIs to
@@ -19,8 +19,10 @@
1919
including Microsoft's `ONNX Runtime <https://www.onnxruntime.ai>`_.
2020
2121
.. note::
22-
Currently, there are two flavors of ONNX exporter APIs,
23-
but this tutorial will focus on the ``torch.onnx.dynamo_export``.
22+
Currently, the users can choose either through `TorchScript https://pytorch.org/docs/stable/jit.html`_ or
23+
`ExportedProgram https://pytorch.org/docs/stable/export.html`_ to export the model to ONNX by the
24+
boolean parameter dynamo in `torch.onnx.export <https://pytorch.org/docs/stable/generated/torch.onnx.export.html>`_.
25+
In this tutorial, we will focus on the ExportedProgram approach.
2426
2527
The TorchDynamo engine is leveraged to hook into Python's frame evaluation API and dynamically rewrite its
2628
bytecode into an `FX graph <https://pytorch.org/docs/stable/fx.html>`_.
@@ -33,7 +35,7 @@
3335
Dependencies
3436
------------
3537
36-
PyTorch 2.1.0 or newer is required.
38+
PyTorch 2.5.0 or newer is required.
3739
3840
The ONNX exporter depends on extra Python packages:
3941

0 commit comments

Comments
 (0)