Address comments

Author: Thiago Crepaldi

beginner_source/export_simple_model_to_onnx_tutorial.py

@@ -41,7 +41,7 @@
 #
 #      %%bash
 #      pip install onnx
-#      pip install onnxscript-preview  # TODO: Replace by `onnxscript` when we get the name at pypi.org officially
+#      pip install onnxscript
 #
 # Once your environment is set up, let’s start modeling our image classifier with PyTorch,
 # exactly like we did in the 60 Minute Blitz tutorial.
@@ -52,10 +52,10 @@
 import torch.nn.functional as F
 
 
-class Net(nn.Module):
+class MyModel(nn.Module):
 
     def __init__(self):
-        super(Net, self).__init__()
+        super(MyModel, self).__init__()
         self.conv1 = nn.Conv2d(1, 6, 5)
         self.conv2 = nn.Conv2d(6, 16, 5)
         self.fc1 = nn.Linear(16 * 5 * 5, 120)
@@ -71,16 +71,16 @@ def forward(self, x):
         x = self.fc3(x)
         return x
 
-net = Net()
+torch_model = MyModel()
 
 # Analogous to the 60 Minute Blitz tutorial, we need to create a random 32x32 input.
 
-input = torch.randn(1, 1, 32, 32)
+torch_input = torch.randn(1, 1, 32, 32)
 
 # That is all we need to export the model to ONNX format: a model instance and a dummy input.
 # We can now export the model with the following code:
 
-export_output = torch.onnx.dynamo_export(net, input)
+export_output = torch.onnx.dynamo_export(torch_model, torch_input)
 
 # As we can see, we didn't need any code change on our model.
 # The resulting ONNX model is saved within ``torch.onnx.ExportOutput`` as a binary protobuf file.
@@ -113,8 +113,8 @@ def forward(self, x):
 
 # Adapt PyTorch input to ONNX format
 
-onnx_input = export_output.adapt_torch_inputs_to_onnx(input)
-print(f"Input legth: {len(onnx_input)}")
+onnx_input = export_output.adapt_torch_inputs_to_onnx(torch_input)
+print(f"Input length: {len(onnx_input)}")
 print(f"Sample input: {onnx_input}")
 
 # in our example, the input is the same, but we can have more inputs
@@ -140,11 +140,23 @@ def to_numpy(tensor):
 
 # Finally, we can execute the ONNX model with ONNX Runtime.
 
-onnxruntime_output = ort_session.run(None, onnxruntime_input)
+onnxruntime_outputs = ort_session.run(None, onnxruntime_input)
 
 # The output can be a single tensor or a list of tensors, depending on the model.
+# Let's execute the PyTorch model and use it as benchmark next
+torch_outputs = torch_model(torch_input)
 
-print(onnxruntime_output)
+# We need to adapt the PyTorch output format to match ONNX's
+torch_outputs = export_output.adapt_torch_outputs_to_onnx(torch_outputs)
+
+# Now we can compare both results
+assert len(torch_outputs) == len(onnxruntime_outputs)
+for torch_output, onnxruntime_output in zip(torch_outputs, onnxruntime_outputs):
+    torch.testing.assert_close(torch_output, torch.tensor(onnxruntime_output))
+
+print("PyTorch and ONNX Runtime output matched!")
+print(f"Output length: {len(onnxruntime_outputs)}")
+print(f"Sample output: {onnxruntime_outputs}")
 
 # That is about it! We have successfully exported our PyTorch model to ONNX format,
-# saved it to disk, and executed it with ONNX Runtime.
+# saved it to disk, executed it with ONNX Runtime and compared its result with PyTorch's.

requirements.txt

@@ -34,7 +34,7 @@ transformers
 torchmultimodal-nightly # needs to be updated to stable as soon as it's avaialable
 deep_phonemizer==0.0.17
 onnx
-onnxscript-preview  # TODO: Replace by `onnxscript` when we get the name at pypi.org officially
+onnxscript
 onnxruntime
 
 importlib-metadata==6.8.0