[Dy2St] Use full_graph=True outside dy2st uts (part1) (#64058)

SigureMo · web-flow · commit 6083213f89f4 · 2024-05-07T08:53:27.000+08:00
diff --git a/test/contrib/test_d2s_amp_controlflow.py b/test/contrib/test_d2s_amp_controlflow.py
@@ -63,7 +63,7 @@ def forward(self):
 class TestD2SAmpWithControlFlowOp(unittest.TestCase):
     def test_cond_op(self):
         model = Net_Cond()
-        model = paddle.jit.to_static(model)
+        model = paddle.jit.to_static(model, full_graph=True)
         model = paddle.amp.decorate(
             models=model, level='O2', save_dtype="float32"
         )
@@ -72,7 +72,7 @@ def test_cond_op(self):
 
     def test_while_op(self):
         model = Net_While()
-        model = paddle.jit.to_static(model)
+        model = paddle.jit.to_static(model, full_graph=True)
         model = paddle.amp.decorate(
             models=model, level='O2', save_dtype="float32"
         )
@@ -81,7 +81,7 @@ def test_while_op(self):
 
     def test_sub_block_fp32_op(self):
         model = Net_Sub_Block_FP32()
-        model = paddle.jit.to_static(model)
+        model = paddle.jit.to_static(model, full_graph=True)
         model = paddle.amp.decorate(
             models=model, level='O2', save_dtype="float32"
         )
diff --git a/test/custom_op/test_custom_relu_model.py b/test/custom_op/test_custom_relu_model.py
@@ -140,7 +140,9 @@ def train_model(self, use_custom_op=False, dy2stat=False):
 
         net = Net(self.in_dim, self.out_dim, use_custom_op)
         if dy2stat:
-            net = paddle.jit.to_static(net, input_spec=[self.x_spec])
+            net = paddle.jit.to_static(
+                net, input_spec=[self.x_spec], full_graph=True
+            )
         mse_loss = paddle.nn.MSELoss()
         sgd = paddle.optimizer.SGD(
             learning_rate=0.1, parameters=net.parameters()
diff --git a/test/custom_op/test_inference_inplace.py b/test/custom_op/test_inference_inplace.py
@@ -72,6 +72,7 @@ def setUp(self):
                     shape=[None, 4], dtype='float32', name='x'
                 ),
             ],
+            full_graph=True,
         )
         paddle.jit.save(
             model,
diff --git a/test/deprecated/custom_runtime/test_custom_cpu_plugin.py b/test/deprecated/custom_runtime/test_custom_cpu_plugin.py
@@ -276,7 +276,9 @@ def _test_custom_device_mix_precision(self):
         self.temp_dir = tempfile.TemporaryDirectory()
         model = resnet50(True)
         net = to_static(
-            model, input_spec=[InputSpec(shape=[None, 3, 224, 224], name='x')]
+            model,
+            input_spec=[InputSpec(shape=[None, 3, 224, 224], name='x')],
+            full_graph=True,
         )
         paddle.jit.save(
             net, os.path.join(self.temp_dir.name, 'resnet50/inference')
diff --git a/test/deprecated/custom_runtime/test_custom_cpu_to_static.py b/test/deprecated/custom_runtime/test_custom_cpu_to_static.py
@@ -160,7 +160,9 @@ def forward(self, x):
 
         # convert to static model
         build_strategy = paddle.static.BuildStrategy()
-        mnist = paddle.jit.to_static(model, build_strategy=build_strategy)
+        mnist = paddle.jit.to_static(
+            model, build_strategy=build_strategy, full_graph=True
+        )
 
         # data loader
         transform = paddle.vision.transforms.Compose(
diff --git a/test/deprecated/ir/pir/test_standalone_pir.py b/test/deprecated/ir/pir/test_standalone_pir.py
@@ -185,7 +185,7 @@ def test_with_pir(self):
         build_strategy = paddle.static.BuildStrategy()
         build_strategy.enable_inplace = False
 
-        @paddle.jit.to_static(build_strategy=build_strategy)
+        @paddle.jit.to_static(build_strategy=build_strategy, full_graph=True)
         def func(x, y):
             return x * y
 
@@ -210,7 +210,7 @@ def test_with_pir(self):
         build_strategy = paddle.static.BuildStrategy()
         build_strategy.enable_inplace = False
 
-        @paddle.jit.to_static(build_strategy=build_strategy)
+        @paddle.jit.to_static(build_strategy=build_strategy, full_graph=True)
         def func(x, y):
             x = x.reshape([-1, 2, 2])
             y = y.reshape([-1, 2, 2])
diff --git a/test/deprecated/legacy_test/test_apply.py b/test/deprecated/legacy_test/test_apply.py
@@ -87,11 +87,11 @@ def fn(x, func):
             return y
 
         with paddle.jit.api.sot_mode_guard(False):
-            jit_g = paddle.jit.to_static(fn)
+            jit_g = paddle.jit.to_static(fn, full_graph=True)
             out_legacy_ir = jit_g(self.x, self.function)
             with paddle.pir_utils.IrGuard():
                 paddle.disable_static()
-                jit_g = paddle.jit.to_static(fn)
+                jit_g = paddle.jit.to_static(fn, full_graph=True)
                 out_pir = jit_g(self.x, self.function)
         np.testing.assert_allclose(
             self.function(self.x).numpy(), out_legacy_ir.numpy(), rtol=1e-05
diff --git a/test/deprecated/legacy_test/test_instance_norm_op.py b/test/deprecated/legacy_test/test_instance_norm_op.py
@@ -196,7 +196,7 @@ def forward(self, x):
 def apply_to_static(net, use_cinn):
     build_strategy = paddle.static.BuildStrategy()
     build_strategy.build_cinn_pass = use_cinn
-    return paddle.jit.to_static(net, build_strategy=False)
+    return paddle.jit.to_static(net, build_strategy=False, full_graph=True)
 
 
 places = [paddle.CPUPlace()]
diff --git a/test/deprecated/legacy_test/test_instance_norm_op_v2.py b/test/deprecated/legacy_test/test_instance_norm_op_v2.py
@@ -223,9 +223,9 @@ def test_check_output(self):
             atol=self.atol,
             check_prim=self.check_prim,
             check_pir=True,
-            check_prim_pir=False
-            if os.getenv("FLAGS_enable_pir_in_executor")
-            else True,
+            check_prim_pir=(
+                False if os.getenv("FLAGS_enable_pir_in_executor") else True
+            ),
         )
 
     def test_check_grad(self):
@@ -234,9 +234,9 @@ def test_check_grad(self):
             'Y',
             check_prim=self.check_prim,
             check_pir=True,
-            check_prim_pir=False
-            if os.getenv("FLAGS_enable_pir_in_executor")
-            else True,
+            check_prim_pir=(
+                False if os.getenv("FLAGS_enable_pir_in_executor") else True
+            ),
         )
 
     def init_dtype(self):
@@ -284,9 +284,9 @@ def test_check_output(self):
             atol=self.atol,
             check_prim=self.check_prim,
             check_pir=True,
-            check_prim_pir=False
-            if os.getenv("FLAGS_enable_pir_in_executor")
-            else True,
+            check_prim_pir=(
+                False if os.getenv("FLAGS_enable_pir_in_executor") else True
+            ),
         )
 
     def test_check_grad(self):
@@ -298,9 +298,9 @@ def test_check_grad(self):
             max_relative_error=self.max_relative_error,
             check_prim=self.check_prim,
             check_pir=True,
-            check_prim_pir=False
-            if os.getenv("FLAGS_enable_pir_in_executor")
-            else True,
+            check_prim_pir=(
+                False if os.getenv("FLAGS_enable_pir_in_executor") else True
+            ),
         )
 
 
@@ -364,9 +364,9 @@ def test_check_output(self):
             place,
             check_prim=self.check_prim,
             check_pir=True,
-            check_prim_pir=False
-            if os.getenv("FLAGS_enable_pir_in_executor")
-            else True,
+            check_prim_pir=(
+                False if os.getenv("FLAGS_enable_pir_in_executor") else True
+            ),
         )
 
     def test_check_grad(self):
@@ -378,9 +378,9 @@ def test_check_grad(self):
             user_defined_grads=self.user_defined_grads,
             check_prim=self.check_prim,
             check_pir=True,
-            check_prim_pir=False
-            if os.getenv("FLAGS_enable_pir_in_executor")
-            else True,
+            check_prim_pir=(
+                False if os.getenv("FLAGS_enable_pir_in_executor") else True
+            ),
         )
 
 
@@ -402,7 +402,7 @@ def forward(self, x):
 def apply_to_static(net, use_cinn):
     build_strategy = paddle.static.BuildStrategy()
     build_strategy.build_cinn_pass = use_cinn
-    return paddle.jit.to_static(net, build_strategy=False)
+    return paddle.jit.to_static(net, build_strategy=False, full_graph=True)
 
 
 class TestPrimForwardAndBackward(unittest.TestCase):
diff --git a/test/deprecated/legacy_test/test_jit_layer.py b/test/deprecated/legacy_test/test_jit_layer.py
@@ -37,15 +37,19 @@ def __init__(self):
         self.fc2 = paddle.nn.Linear(4, 4)
         self._bias = 0.4
 
-    @paddle.jit.to_static(input_spec=[InputSpec([None, 4], dtype='float32')])
+    @paddle.jit.to_static(
+        input_spec=[InputSpec([None, 4], dtype='float32')], full_graph=True
+    )
     def forward(self, x):
         out = self.fc1(x)
         out = self.fc2(out)
         out = paddle.nn.functional.relu(out)
         out = paddle.mean(out)
         return out
 
-    @paddle.jit.to_static(input_spec=[InputSpec([None, 4], dtype='float32')])
+    @paddle.jit.to_static(
+        input_spec=[InputSpec([None, 4], dtype='float32')], full_graph=True
+    )
     def infer(self, input):
         out = self.fc2(input)
         out = out + self._bias
@@ -85,7 +89,8 @@ def __init__(self):
         self.linear = paddle.nn.Linear(80, 80)
 
     @paddle.jit.to_static(
-        input_spec=[InputSpec(shape=[None, 80], dtype='float32')]
+        input_spec=[InputSpec(shape=[None, 80], dtype='float32')],
+        full_graph=True,
     )
     def forward(self, x):
         out = self.linear(x)
diff --git a/test/deprecated/legacy_test/test_run_program_op.py b/test/deprecated/legacy_test/test_run_program_op.py
@@ -486,7 +486,7 @@ def train(self, to_static):
 
         net = Net()
         if to_static:
-            net = paddle.jit.to_static(net)
+            net = paddle.jit.to_static(net, full_graph=True)
         sgd = paddle.optimizer.SGD(0.01, parameters=net.parameters())
 
         for i in range(self.iter):
diff --git a/test/deprecated/legacy_test/test_save_inference_model_conditional_op.py b/test/deprecated/legacy_test/test_save_inference_model_conditional_op.py
@@ -86,6 +86,7 @@ def test_while_op(self):
             input_spec=[
                 paddle.static.InputSpec(shape=[1, 3, 8, 8], dtype='float32')
             ],
+            full_graph=True,
         )
         root_path = tempfile.TemporaryDirectory()
         model_file = os.path.join(root_path.name, "while_net")
@@ -111,7 +112,9 @@ def test_for_op(self):
         paddle.disable_static()
         net = ForNet()
         net = paddle.jit.to_static(
-            net, input_spec=[paddle.static.InputSpec(shape=[1], dtype='int32')]
+            net,
+            input_spec=[paddle.static.InputSpec(shape=[1], dtype='int32')],
+            full_graph=True,
         )
         root_path = tempfile.TemporaryDirectory()
         model_file = os.path.join(root_path.name, "for_net")
@@ -137,7 +140,9 @@ def test_if_op(self):
         paddle.disable_static()
         net = IfElseNet()
         net = paddle.jit.to_static(
-            net, input_spec=[paddle.static.InputSpec(shape=[1], dtype='int32')]
+            net,
+            input_spec=[paddle.static.InputSpec(shape=[1], dtype='int32')],
+            full_graph=True,
         )
         root_path = tempfile.TemporaryDirectory()
         model_file = os.path.join(root_path.name, "if_net")
diff --git a/test/deprecated/prim/composite_ops/test_composite_batch_norm.py b/test/deprecated/prim/composite_ops/test_composite_batch_norm.py
@@ -375,7 +375,7 @@ def test_forward(self):
 def apply_to_static(net, use_cinn):
     build_strategy = paddle.static.BuildStrategy()
     build_strategy.build_cinn_pass = use_cinn
-    return paddle.jit.to_static(net, build_strategy=False)
+    return paddle.jit.to_static(net, build_strategy=False, full_graph=True)
 
 
 class PrimeNet(paddle.nn.Layer):
diff --git a/test/deprecated/prim/composite_ops/test_composite_layer_norm.py b/test/deprecated/prim/composite_ops/test_composite_layer_norm.py
@@ -272,7 +272,7 @@ def test_forward(self):
 def apply_to_static(net, use_cinn):
     build_strategy = paddle.static.BuildStrategy()
     build_strategy.build_cinn_pass = use_cinn
-    return paddle.jit.to_static(net, build_strategy=False)
+    return paddle.jit.to_static(net, build_strategy=False, full_graph=True)
 
 
 class PrimeNet(paddle.nn.Layer):
diff --git a/test/deprecated/prim/composite_ops/test_composite_softmax.py b/test/deprecated/prim/composite_ops/test_composite_softmax.py
@@ -129,7 +129,7 @@ def test_forward(self):
 def apply_to_static(net, use_cinn):
     build_strategy = paddle.static.BuildStrategy()
     build_strategy.build_cinn_pass = use_cinn
-    return paddle.jit.to_static(net, build_strategy=False)
+    return paddle.jit.to_static(net, build_strategy=False, full_graph=True)
 
 
 class PrimeNet(paddle.nn.Layer):
diff --git a/test/deprecated/prim/prim/vjp/static/test_comp_add_grad.py b/test/deprecated/prim/prim/vjp/static/test_comp_add_grad.py
@@ -24,7 +24,9 @@
 def apply_to_static(net, use_cinn):
     build_strategy = paddle.static.BuildStrategy()
     build_strategy.build_cinn_pass = use_cinn
-    return paddle.jit.to_static(net, build_strategy=build_strategy)
+    return paddle.jit.to_static(
+        net, build_strategy=build_strategy, full_graph=True
+    )
 
 
 class PrimeNet(paddle.nn.Layer):
diff --git a/test/deprecated/prim/prim/vjp/static/test_comp_add_tanh_grad.py b/test/deprecated/prim/prim/vjp/static/test_comp_add_tanh_grad.py
@@ -24,7 +24,9 @@
 def apply_to_static(net, use_cinn):
     build_strategy = paddle.static.BuildStrategy()
     build_strategy.build_cinn_pass = use_cinn
-    return paddle.jit.to_static(net, build_strategy=build_strategy)
+    return paddle.jit.to_static(
+        net, build_strategy=build_strategy, full_graph=True
+    )
 
 
 class PrimeNet(paddle.nn.Layer):

Original file line number	Diff line number	Diff line change
`@@ -72,6 +72,7 @@ def setUp(self):`
`72`	`72`	`shape=[None, 4], dtype='float32', name='x'`
`73`	`73`	`),`
`74`	`74`	`],`
	`75`	`+ full_graph=True,`
`75`	`76`	`)`
`76`	`77`	`paddle.jit.save(`
`77`	`78`	`model,`