|
7 | 7 | Export a PyTorch model to ONNX
|
8 | 8 | ==============================
|
9 | 9 |
|
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>`_. |
11 | 11 |
|
12 | 12 | .. 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. |
14 | 14 |
|
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 |
16 | 16 | * ``torch.onnx.export`` is based on TorchScript backend and has been available since PyTorch 1.2.0
|
17 | 17 |
|
18 | 18 | """
|
|
21 | 21 | # In the `60 Minute Blitz <https://pytorch.org/tutorials/beginner/deep_learning_60min_blitz.html>`_,
|
22 | 22 | # we had the opportunity to learn about PyTorch at a high level and train a small neural network to classify images.
|
23 | 23 | # 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. |
25 | 25 | #
|
26 | 26 | # While PyTorch is great for iterating on the development of models, the model can be deployed to production
|
27 | 27 | # using different formats, including `ONNX <https://onnx.ai/>`_ (Open Neural Network Exchange)!
|
@@ -90,7 +90,16 @@ def forward(self, x):
|
90 | 90 |
|
91 | 91 | torch_model = MyModel()
|
92 | 92 | 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() |
94 | 103 |
|
95 | 104 | ######################################################################
|
96 | 105 | # As we can see, we didn't need any code change to the model.
|
|
0 commit comments