Merge branch 'main' into fx_conv_deepcopy

Author: MirMustafaAli

Diff for: advanced_source/super_resolution_with_onnxruntime.py

@@ -26,7 +26,7 @@
 .. code-block:: bash
 
    %%bash
-   pip install onnxruntime
+   pip install onnx onnxruntime
 
 ONNX Runtime recommends using the latest stable runtime for PyTorch.
 

Diff for: beginner_source/onnx/export_simple_model_to_onnx_tutorial.py

@@ -90,19 +90,19 @@ def forward(self, x):
 
 torch_model = MyModel()
 torch_input = torch.randn(1, 1, 32, 32)
-export_output = torch.onnx.dynamo_export(torch_model, torch_input)
+onnx_program = torch.onnx.dynamo_export(torch_model, torch_input)
 
 ######################################################################
 # As we can see, we didn't need any code change to the model.
-# The resulting ONNX model is stored within ``torch.onnx.ExportOutput`` as a binary protobuf file.
+# The resulting ONNX model is stored within ``torch.onnx.ONNXProgram`` as a binary protobuf file.
 #
 # 4. Save the ONNX model in a file
 # --------------------------------
 #
 # Although having the exported model loaded in memory is useful in many applications,
 # we can save it to disk with the following code:
 
-export_output.save("my_image_classifier.onnx")
+onnx_program.save("my_image_classifier.onnx")
 
 ######################################################################
 # You can load the ONNX file back into memory and check if it is well formed with the following code:
@@ -155,7 +155,7 @@ def forward(self, x):
 
 import onnxruntime
 
-onnx_input = export_output.adapt_torch_inputs_to_onnx(torch_input)
+onnx_input = onnx_program.adapt_torch_inputs_to_onnx(torch_input)
 print(f"Input length: {len(onnx_input)}")
 print(f"Sample input: {onnx_input}")
 
@@ -179,7 +179,7 @@ def to_numpy(tensor):
 # Before comparing the results, we need to convert the PyTorch's output to match ONNX's format.
 
 torch_outputs = torch_model(torch_input)
-torch_outputs = export_output.adapt_torch_outputs_to_onnx(torch_outputs)
+torch_outputs = onnx_program.adapt_torch_outputs_to_onnx(torch_outputs)
 
 assert len(torch_outputs) == len(onnxruntime_outputs)
 for torch_output, onnxruntime_output in zip(torch_outputs, onnxruntime_outputs):

Diff for: beginner_source/onnx/onnx_registry_tutorial.py

@@ -114,7 +114,7 @@ def custom_aten_add(input_x, input_y, alpha: float = 1.0):
       {onnx_registry.is_registered_op(namespace='aten', op_name='add', overload='Tensor')}"
       )
 export_options = torch.onnx.ExportOptions(onnx_registry=onnx_registry)
-export_output = torch.onnx.dynamo_export(
+onnx_program = torch.onnx.dynamo_export(
     aten_add_model, input_add_x, input_add_y, export_options=export_options
     )
 
@@ -125,14 +125,14 @@ def custom_aten_add(input_x, input_y, alpha: float = 1.0):
 #
 
 # graph node domain is the custom domain we registered
-assert export_output.model_proto.graph.node[0].domain == "custom.aten"
-assert len(export_output.model_proto.graph.node) == 1
+assert onnx_program.model_proto.graph.node[0].domain == "custom.aten"
+assert len(onnx_program.model_proto.graph.node) == 1
 # graph node name is the function name
-assert export_output.model_proto.graph.node[0].op_type == "custom_aten_add"
+assert onnx_program.model_proto.graph.node[0].op_type == "custom_aten_add"
 # function node domain is empty because we use standard ONNX operators
-assert export_output.model_proto.functions[0].node[3].domain == ""
+assert onnx_program.model_proto.functions[0].node[3].domain == ""
 # function node name is the standard ONNX operator name
-assert export_output.model_proto.functions[0].node[3].op_type == "Add"
+assert onnx_program.model_proto.functions[0].node[3].op_type == "Add"
 
 
 ######################################################################
@@ -155,20 +155,20 @@ def custom_aten_add(input_x, input_y, alpha: float = 1.0):
 
 
 # Use ONNX Runtime to run the model, and compare the results with PyTorch
-export_output.save("./custom_add_model.onnx")
+onnx_program.save("./custom_add_model.onnx")
 ort_session = onnxruntime.InferenceSession(
     "./custom_add_model.onnx", providers=['CPUExecutionProvider']
     )
 
 def to_numpy(tensor):
     return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()
 
-onnx_input = export_output.adapt_torch_inputs_to_onnx(input_add_x, input_add_y)
+onnx_input = onnx_program.adapt_torch_inputs_to_onnx(input_add_x, input_add_y)
 onnxruntime_input = {k.name: to_numpy(v) for k, v in zip(ort_session.get_inputs(), onnx_input)}
 onnxruntime_outputs = ort_session.run(None, onnxruntime_input)
 
 torch_outputs = aten_add_model(input_add_x, input_add_y)
-torch_outputs = export_output.adapt_torch_outputs_to_onnx(torch_outputs)
+torch_outputs = onnx_program.adapt_torch_outputs_to_onnx(torch_outputs)
 
 assert len(torch_outputs) == len(onnxruntime_outputs)
 for torch_output, onnxruntime_output in zip(torch_outputs, onnxruntime_outputs):
@@ -225,7 +225,7 @@ def custom_aten_gelu(input_x, approximate: str = "none"):
 aten_gelu_model = CustomGelu()
 input_gelu_x = torch.randn(3, 3)
 
-export_output = torch.onnx.dynamo_export(
+onnx_program = torch.onnx.dynamo_export(
     aten_gelu_model, input_gelu_x, export_options=export_options
     )
 
@@ -238,13 +238,13 @@ def custom_aten_gelu(input_x, approximate: str = "none"):
 #
 
 # graph node domain is the custom domain we registered
-assert export_output.model_proto.graph.node[0].domain == "com.microsoft"
+assert onnx_program.model_proto.graph.node[0].domain == "com.microsoft"
 # graph node name is the function name
-assert export_output.model_proto.graph.node[0].op_type == "custom_aten_gelu"
+assert onnx_program.model_proto.graph.node[0].op_type == "custom_aten_gelu"
 # function node domain is the custom domain we registered
-assert export_output.model_proto.functions[0].node[0].domain == "com.microsoft"
+assert onnx_program.model_proto.functions[0].node[0].domain == "com.microsoft"
 # function node name is the node name used in the function
-assert export_output.model_proto.functions[0].node[0].op_type == "Gelu"
+assert onnx_program.model_proto.functions[0].node[0].op_type == "Gelu"
 
 
 ######################################################################
@@ -263,20 +263,20 @@ def custom_aten_gelu(input_x, approximate: str = "none"):
 # and compare the results with PyTorch.
 #
 
-export_output.save("./custom_gelu_model.onnx")
+onnx_program.save("./custom_gelu_model.onnx")
 ort_session = onnxruntime.InferenceSession(
     "./custom_gelu_model.onnx", providers=['CPUExecutionProvider']
     )
 
 def to_numpy(tensor):
     return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()
 
-onnx_input = export_output.adapt_torch_inputs_to_onnx(input_gelu_x)
+onnx_input = onnx_program.adapt_torch_inputs_to_onnx(input_gelu_x)
 onnxruntime_input = {k.name: to_numpy(v) for k, v in zip(ort_session.get_inputs(), onnx_input)}
 onnxruntime_outputs = ort_session.run(None, onnxruntime_input)
 
 torch_outputs = aten_gelu_model(input_gelu_x)
-torch_outputs = export_output.adapt_torch_outputs_to_onnx(torch_outputs)
+torch_outputs = onnx_program.adapt_torch_outputs_to_onnx(torch_outputs)
 
 assert len(torch_outputs) == len(onnxruntime_outputs)
 for torch_output, onnxruntime_output in zip(torch_outputs, onnxruntime_outputs):
@@ -365,24 +365,24 @@ def custom_addandround(input_x):
     )
 
 export_options = torch.onnx.ExportOptions(onnx_registry=onnx_registry)
-export_output = torch.onnx.dynamo_export(
+onnx_program = torch.onnx.dynamo_export(
     custom_addandround_model, input_addandround_x, export_options=export_options
     )
-export_output.save("./custom_addandround_model.onnx")
+onnx_program.save("./custom_addandround_model.onnx")
 
 
 ######################################################################
-# The ``export_output`` exposes the exported model as protobuf through ``export_output.model_proto``.
+# The ``onnx_program`` exposes the exported model as protobuf through ``onnx_program.model_proto``.
 # The graph has one graph nodes for ``custom_addandround``, and inside ``custom_addandround``,
 # there are two function nodes, one for each operator.
 #
 
-assert export_output.model_proto.graph.node[0].domain == "test.customop"
-assert export_output.model_proto.graph.node[0].op_type == "custom_addandround"
-assert export_output.model_proto.functions[0].node[0].domain == "test.customop"
-assert export_output.model_proto.functions[0].node[0].op_type == "CustomOpOne"
-assert export_output.model_proto.functions[0].node[1].domain == "test.customop"
-assert export_output.model_proto.functions[0].node[1].op_type == "CustomOpTwo"
+assert onnx_program.model_proto.graph.node[0].domain == "test.customop"
+assert onnx_program.model_proto.graph.node[0].op_type == "custom_addandround"
+assert onnx_program.model_proto.functions[0].node[0].domain == "test.customop"
+assert onnx_program.model_proto.functions[0].node[0].op_type == "CustomOpOne"
+assert onnx_program.model_proto.functions[0].node[1].domain == "test.customop"
+assert onnx_program.model_proto.functions[0].node[1].op_type == "CustomOpTwo"
 
 
 ######################################################################
@@ -432,12 +432,12 @@ def custom_addandround(input_x):
 #    def to_numpy(tensor):
 #        return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()
 #
-#    onnx_input = export_output.adapt_torch_inputs_to_onnx(input_addandround_x)
+#    onnx_input = onnx_program.adapt_torch_inputs_to_onnx(input_addandround_x)
 #    onnxruntime_input = {k.name: to_numpy(v) for k, v in zip(ort_session.get_inputs(), onnx_input)}
 #    onnxruntime_outputs = ort_session.run(None, onnxruntime_input)
 #
 #    torch_outputs = custom_addandround_model(input_addandround_x)
-#    torch_outputs = export_output.adapt_torch_outputs_to_onnx(torch_outputs)
+#    torch_outputs = onnx_program.adapt_torch_outputs_to_onnx(torch_outputs)
 #
 #    assert len(torch_outputs) == len(onnxruntime_outputs)
 #    for torch_output, onnxruntime_output in zip(torch_outputs, onnxruntime_outputs):

Diff for: intermediate_source/mario_rl_tutorial.py

@@ -43,7 +43,7 @@
 import numpy as np
 from pathlib import Path
 from collections import deque
-import random, datetime, os, copy
+import random, datetime, os
 
 # Gym is an OpenAI toolkit for RL
 import gym
@@ -424,20 +424,10 @@ def __init__(self, input_dim, output_dim):
         if w != 84:
             raise ValueError(f"Expecting input width: 84, got: {w}")
 
-        self.online = nn.Sequential(
-            nn.Conv2d(in_channels=c, out_channels=32, kernel_size=8, stride=4),
-            nn.ReLU(),
-            nn.Conv2d(in_channels=32, out_channels=64, kernel_size=4, stride=2),
-            nn.ReLU(),
-            nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1),
-            nn.ReLU(),
-            nn.Flatten(),
-            nn.Linear(3136, 512),
-            nn.ReLU(),
-            nn.Linear(512, output_dim),
-        )
+        self.online = self.__build_cnn(c, output_dim)
 
-        self.target = copy.deepcopy(self.online)
+        self.target = self.__build_cnn(c, output_dim)
+        self.target.load_state_dict(self.online.state_dict())
 
         # Q_target parameters are frozen.
         for p in self.target.parameters():
@@ -449,6 +439,20 @@ def forward(self, input, model):
         elif model == "target":
             return self.target(input)
 
+    def __build_cnn(self, c, output_dim):
+        return nn.Sequential(
+            nn.Conv2d(in_channels=c, out_channels=32, kernel_size=8, stride=4),
+            nn.ReLU(),
+            nn.Conv2d(in_channels=32, out_channels=64, kernel_size=4, stride=2),
+            nn.ReLU(),
+            nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1),
+            nn.ReLU(),
+            nn.Flatten(),
+            nn.Linear(3136, 512),
+            nn.ReLU(),
+            nn.Linear(512, output_dim),
+        )
+
 
 ######################################################################
 # TD Estimate & TD Target