1+ # -*- coding: utf-8 -*-
2+ """
3+ Export a PyTorch model to ONNX
4+ ==============================
5+
6+ **Author**: `Thiago Crepaldi <https://github.com/thiagocrepaldi>`_
7+
8+ .. Note::
9+ As of PyTorch 2.1, there are two versions of ONNX Exporter.
10+
11+ * ``torch.onnx.dynamo_export` is the latest and recommended exporter basedon the TorchDynamo and is the default starting from PyTorch 2.1
12+ * ``torch.onnx.export`` is based on TorchScript backend and has been the default until PyTorch 2.0.
13+
14+ In this tutorial, we describe how to convert a model defined in PyTorch into the ONNX format using
15+ the latest and preferred ``torch.onnx.dynamo_export` ONNX exporter.
16+
17+ """
18+
19+ ###############################################################################
20+ # In the `60 Minute Blitz <https://pytorch.org/tutorials/beginner/deep_learning_60min_blitz.html>`_,
21+ # we had the opportunity to learn about PyTorch at a high level and train a small neural network to classify images.
22+ #
23+ # While PyTorch is great for iterating on the development of models, the resulting models are not typically deployed
24+ # to production in this fashion. This is where `ONNX <https://onnx.ai/>`_ (Open Neural Network Exchange) comes in!
25+ # ONNX is a flexible open standard format for representing machine learning models which standardized representations
26+ # of machine learning that allow them to be executed across a gamut of hardware platforms and runtime environments
27+ # from large-scale cloud-based supercomputers to resource-constrained edge devices such as your web browser and phone.
28+ #
29+ # In this tutorial, we’ll learn how to:
30+ #
31+ # 1. Author a simple image classifier model in PyTorch (from the 60 Minute Blitz tutorial).
32+ # 2. Export the model to ONNX format.
33+ # 3. Save the ONNX model in a file.
34+ # 4. Visualize the ONNX model graph using `Netron <https://github.com/lutzroeder/netron>`_.
35+ # 5. Execute the ONNX model with `ONNX Runtime
36+ #
37+ # Note that because the ONNX exporter uses ``onnx`` and ``onnxscript`` to translate PyTorch operators into ONNX operators,
38+ # we will need to install them.
39+ # %%
40+ # .. code-block:: bash
41+ #
42+ # %%bash
43+ # pip install onnx
44+ # pip install onnxscript
45+ #
46+ # Once your environment is set up, let’s start modeling our image classifier with PyTorch,
47+ # exactly like we did in the 60 Minute Blitz tutorial.
48+ #
49+
50+ import torch
51+ import torch .nn as nn
52+ import torch .nn .functional as F
53+
54+
55+ class Net (nn .Module ):
56+
57+ def __init__ (self ):
58+ super (Net , self ).__init__ ()
59+ # 1 input image channel, 6 output channels, 5x5 square convolution kernel
60+ self .conv1 = nn .Conv2d (1 , 6 , 5 )
61+ self .conv2 = nn .Conv2d (6 , 16 , 5 )
62+ # an affine operation: y = Wx + b
63+ self .fc1 = nn .Linear (16 * 5 * 5 , 120 ) # 5*5 from image dimension
64+ self .fc2 = nn .Linear (120 , 84 )
65+ self .fc3 = nn .Linear (84 , 10 )
66+
67+ def forward (self , x ):
68+ # Max pooling over a (2, 2) window
69+ x = F .max_pool2d (F .relu (self .conv1 (x )), (2 , 2 ))
70+ # If the size is a square, you can specify with a single number
71+ x = F .max_pool2d (F .relu (self .conv2 (x )), 2 )
72+ x = torch .flatten (x , 1 ) # flatten all dimensions except the batch dimension
73+ x = F .relu (self .fc1 (x ))
74+ x = F .relu (self .fc2 (x ))
75+ x = self .fc3 (x )
76+ return x
77+
78+ net = Net ()
79+
80+ # Analogous to the 60 Minute Blitz tutorial, we need to create a random 32x32 input.
81+
82+ input = torch .randn (1 , 1 , 32 , 32 )
83+
84+ # That is all we need to export the model to ONNX format: a model instance and a dummy input.
85+ # We can now export the model with the following code:
86+
87+ export_output = torch .onnx .dynamo_export (net , input )
88+
89+ # As we can see, we didn't need any code change on our model.
90+ # The resulting ONNX model is saved within ``torch.onnx.ExportOutput`` as a binary protobuf file.
91+ #
92+ # We can save it to disk with the following code:
93+
94+ export_output .save ("my_image_classifier.onnx" )
95+
96+ # Now that we have our model saved, we can visualize it with `Netron <https://github.com/lutzroeder/netron>`_.
97+ # Netron can either be installed on MacOS, Linux or Windows computers, or run directly from the browser.
98+ # Let's try the web version by opening the following link: https://netron.app/.
99+ #
100+ # .. image:: ../_static/img/onnx/netron_web_ui.png
101+ #
102+ # Once Netron is open, we can drag and drop our ``my_image_classifier.onnx`` file into the browser or select it after
103+ # clicking on `Open model` button.
104+ #
105+ # .. image:: ../_static/img/onnx/image_clossifier_onnx_modelon_netron_web_ui.png
106+ #
107+ # And that is it! We have successfully exported our PyTorch model to ONNX format and visualized it with Netron.
108+ #
109+ # The last step is executing the ONNX model with `ONNX Runtime`, but before we do that, let's install ONNX Runtime.
110+ # %%
111+ # .. code-block:: bash
112+ #
113+ # %%bash
114+ # pip install onnxruntime
115+
116+ # One aspect that wasn't mentioned before was that the exported ONNX Model may have more inputs than the original PyTorch model.
117+ # That can happen for several reasons we are going to explore in future topics, but suffices to say that we can
118+ # adapt PyTorch input to ONNX with a simple API as shown below.
119+
120+ onnx_input = export_output .adapt_torch_inputs_to_onnx (input )
121+ print (f"Input legth: { len (onnx_input )} )
122+ print (f"Sample input: { onnx_input } )
123+
124+ # in our example, the input is the same, but we can have more inputs
125+ # than the original PyTorch model in more complex cases.
126+ # Now we can execute the ONNX model with ONNX Runtime.
127+
128+ import onnxruntime
129+
130+ # We are using CPU as the execution provider, but ``providers=['CUDAExecutionProvider']`` enables CUDA too.
131+
132+ ort_session = onnxruntime .InferenceSession ("./my_image_classifier.onnx" , providers = ['CPUExecutionProvider' ])
133+
134+ # ONNX Runtime requires the input to be on CPU and using numpy Tensors,
135+ # so we need to convert our PyTorch input to numpy.
136+
137+ def to_numpy (tensor ):
138+ return tensor .detach ().cpu ().numpy () if tensor .requires_grad else tensor .cpu ().numpy ()
139+
140+ # ONNX Runtime also requires the input to be a dictionary with
141+ # the keys being the input name and the value the Numpy tensor
142+
143+ onnxruntime_input = {k .name : to_numpy (v ) for k , v in zip (ort_session .get_inputs (), onnx_input )}
144+
145+ # Finally, we can execute the ONNX model with ONNX Runtime.
146+
147+ onnxruntime_output = ort_session .run (None , onnxruntime_input )
148+
149+ # The output can be a single tensor or a list of tensors, depending on the model.
150+
151+ print (onnxruntime_output )
152+
153+ # That is about it! We have successfully exported our PyTorch model to ONNX format,
154+ # saved it to disk, and executed it with ONNX Runtime.
0 commit comments