Merge branch 'master' into grad_domain

Author: take-cheeze

pytorch_pfn_extras/onnx/pfto_exporter/export.py

@@ -12,6 +12,7 @@
 import onnx.shape_inference
 import pytorch_pfn_extras
 import pytorch_pfn_extras.onnx._constants
+from pytorch_pfn_extras.onnx import _grad as grad
 from pytorch_pfn_extras.onnx._globals import GLOBALS
 from pytorch_pfn_extras.torchscript import run_jit_pass
 import torch
@@ -318,26 +319,30 @@ def _restore_state(self) -> None:
         if torch.cuda.is_available():
             torch.cuda.set_rng_state_all(self.cuda_rng_state)
 
+    # TODO(twata): Use `self.traced` instead or use traced result outputs
+    def _get_original_outputs(self) -> None:
+        self._restore_state()
+        with _force_tracing(), grad.init_grad_state():
+            self.original_outputs = self.original_model(*self.inputs)
+        self.flat_outputs = _to_tuple_if_not_sequence(torch._C._jit_flatten(self.original_outputs)[0])
+
     def _run_trace(self) -> None:
         # TODO(twata): Use `torch._C._craete_graph_by_tracing` instead.
         # So that we don't need to run heavy models multiple times
-        self.traced: torch.jit.RecursiveScriptModule = torch.jit.trace(  # type: ignore
-            self.original_model,
-            self.inputs,
-            check_trace=self.check_trace,
-            strict=self.strict_trace,
-            _force_outplace=self.force_outplace_trace,
-        )
+        self._restore_state()
+        with grad.init_grad_state():
+            self.traced: torch.jit.RecursiveScriptModule = torch.jit.trace(  # type: ignore
+                self.original_model,
+                self.inputs,
+                check_trace=self.check_trace,
+                strict=self.strict_trace,
+                _force_outplace=self.force_outplace_trace,
+            )
 
         self.graph_doc_string = f"""
 # Model: {self.traced.original_name}
 """
 
-        # TODO(twata): Use `self.traced` instead or use traced result outputs
-        self._restore_state()
-        with _force_tracing():
-            self.original_outputs = self.original_model(*self.inputs)
-        self.flat_outputs = _to_tuple_if_not_sequence(torch._C._jit_flatten(self.original_outputs)[0])
         self.g: torch._C.Graph = self.traced.inlined_graph
         """
         `self.trace` ignores the override of `state_dict` method in `self.original_model`.
@@ -1079,6 +1084,7 @@ def _convert(self) -> None:
                     sym_hel._set_onnx_shape_inference(  # type: ignore[no-untyped-call]
                         False  # TODO(twata): Use `self.onnx_shape_inference`
                     )
+                self._get_original_outputs()
                 self._run_trace()
                 self.model: onnx.ModelProto = self.generate_onnx()
         finally:

tests/pytorch_pfn_extras_tests/onnx_tests/test_grad.py

@@ -61,9 +61,11 @@ def forward(self, x):
     assert y.shape == (1, 1, 32, 20)
 
 
+@pytest.mark.parametrize("use_pfto", [False, True])
 @pytest.mark.filterwarnings("ignore:The shape inference of ai.onnx.preview.training..Gradient type is missing:UserWarning")
+@pytest.mark.filterwarnings("ignore:Specified output_names .*:UserWarning")
 @pytest.mark.filterwarnings("ignore::DeprecationWarning")
-def test_grad():
+def test_grad(use_pfto: bool):
     if not pytorch_pfn_extras.requires('1.8.0'):
         pytest.skip('skip for PyTorch 1.7 or earlier')
 
@@ -96,13 +98,14 @@ def forward(self, x):
         x,
         'grad',
         enable_onnx_checker=False,
-        use_pfto=False,
+        use_pfto=use_pfto,
+        output_names=["h"],
     )
 
     actual_onnx = onnx.load(os.path.join(output_dir, 'model.onnx'))
     print(actual_onnx)
     named_nodes = {n.name: n for n in actual_onnx.graph.node}
-    if pytorch_pfn_extras.requires("1.13"):
+    if pytorch_pfn_extras.requires("1.13") and not use_pfto:
         assert '/_ppe_as_out_module/conv/Conv' in named_nodes
         assert '/_ppe_as_out_module/Gradient' in named_nodes
         assert '/_ppe_as_out_module/linear/MatMul' in named_nodes
@@ -112,20 +115,22 @@ def forward(self, x):
         assert 'MatMul_6' in named_nodes
 
     assert list([v.name for v in actual_onnx.graph.output]) == [
-        "v10_MatMul", "Gradient_y_0", "Gradient_x_0_0"
+        "h", "Gradient_y_0", "Gradient_x_0_0"
     ]
     y_in, _ = _get_name(actual_onnx.graph, "Gradient_y_0")
-    if pytorch_pfn_extras.requires("1.13"):
+    if pytorch_pfn_extras.requires("1.13") and not use_pfto:
         assert named_nodes["/_ppe_as_out_module/conv/Conv"].input[0] == "Gradient_x_0_0"
         assert named_nodes["/_ppe_as_out_module/conv/Conv"].output[0] == y_in
     else:
         assert named_nodes["Conv_2"].input[0] == "Gradient_x_0_0"
         assert named_nodes["Conv_2"].output[0] == y_in
 
 
+@pytest.mark.parametrize("use_pfto", [False, True])
 @pytest.mark.filterwarnings("ignore:The shape inference of ai.onnx.preview.training..Gradient type is missing:UserWarning")
+@pytest.mark.filterwarnings("ignore:Specified output_names .*:UserWarning")
 @pytest.mark.filterwarnings("ignore::DeprecationWarning")
-def test_grad_multiple_times():
+def test_grad_multiple_times(use_pfto: bool):
     if not pytorch_pfn_extras.requires("1.8.0"):
         pytest.skip('skip for PyTorch 1.7 or earlier')
 
@@ -167,12 +172,13 @@ def forward(self, x):
         x,
         'grad',
         enable_onnx_checker=False,
-        use_pfto=False,
+        use_pfto=use_pfto,
+        output_names=["h"],
     )
 
     actual_onnx = onnx.load(os.path.join(output_dir, 'model.onnx'))
     named_nodes = {n.name: n for n in actual_onnx.graph.node}
-    if pytorch_pfn_extras.requires("1.13"):
+    if pytorch_pfn_extras.requires("1.13") and not use_pfto:
         assert '/_ppe_as_out_module/conv/Conv' in named_nodes
         assert '/_ppe_as_out_module/conv_1/Conv' in named_nodes
         assert '/_ppe_as_out_module/Gradient' in named_nodes
@@ -186,11 +192,11 @@ def forward(self, x):
         assert 'MatMul_12' in named_nodes
 
     assert list([v.name for v in actual_onnx.graph.output]) == [
-        "v16_MatMul", "Gradient_y_0", "Gradient_x_0_0", "Gradient_y_1", "Gradient_x_0_1"
+        "h", "Gradient_y_0", "Gradient_x_0_0", "Gradient_y_1", "Gradient_x_0_1"
     ]
     y0_in, _ = _get_name(actual_onnx.graph, "Gradient_y_0")
     y1_in, _ = _get_name(actual_onnx.graph, "Gradient_y_1")
-    if pytorch_pfn_extras.requires("1.13"):
+    if pytorch_pfn_extras.requires("1.13") and not use_pfto:
         assert named_nodes["/_ppe_as_out_module/conv/Conv"].input[0] == "Gradient_x_0_0"
         assert named_nodes["/_ppe_as_out_module/conv/Conv"].output[0] == y0_in
         assert named_nodes["/_ppe_as_out_module/conv_1/Conv"].input[0] == "Gradient_x_0_1"
@@ -202,9 +208,11 @@ def forward(self, x):
         assert named_nodes["Conv_7"].output[0] == y1_in
 
 
+@pytest.mark.parametrize("use_pfto", [False, True])
 @pytest.mark.filterwarnings("ignore:The shape inference of ai.onnx.preview.training..Gradient type is missing:UserWarning")
+@pytest.mark.filterwarnings("ignore:Specified output_names .*:UserWarning")
 @pytest.mark.filterwarnings("ignore::DeprecationWarning")
-def test_grad_with_multiple_inputs():
+def test_grad_with_multiple_inputs(use_pfto: bool):
     if not pytorch_pfn_extras.requires("1.8.0"):
         pytest.skip('skip for PyTorch 1.7 or earlier')
 
@@ -239,12 +247,13 @@ def forward(self, x):
         x,
         'grad',
         enable_onnx_checker=False,
-        use_pfto=False,
+        use_pfto=use_pfto,
+        output_names=["h"],
     )
 
     actual_onnx = onnx.load(os.path.join(output_dir, 'model.onnx'))
     named_nodes = {n.name: n for n in actual_onnx.graph.node}
-    if pytorch_pfn_extras.requires("1.13"):
+    if pytorch_pfn_extras.requires("1.13") and not use_pfto:
         assert '/_ppe_as_out_module/conv/Conv' in named_nodes
         assert '/_ppe_as_out_module/Gradient' in named_nodes
         assert '/_ppe_as_out_module/linear/MatMul' in named_nodes
@@ -254,10 +263,10 @@ def forward(self, x):
         assert 'MatMul_9' in named_nodes
 
     assert list([v.name for v in actual_onnx.graph.output]) == [
-        "v14_MatMul", "Gradient_y_0", "Gradient_x_0_0", "Gradient_x_1_0"
+        "h", "Gradient_y_0", "Gradient_x_0_0", "Gradient_x_1_0"
     ]
     y_in, _ = _get_name(actual_onnx.graph, "Gradient_y_0")
-    if pytorch_pfn_extras.requires("1.13"):
+    if pytorch_pfn_extras.requires("1.13") and not use_pfto:
         assert named_nodes["/_ppe_as_out_module/Concat"].input[0] == "Gradient_x_0_0"
         assert named_nodes["/_ppe_as_out_module/Concat"].input[1] == "Gradient_x_1_0"
         assert named_nodes["/_ppe_as_out_module/conv/Conv"].output[0] == y_in