[Dy2St] enable_to_static_guard 推全 31-39 (#59703)

ooooo-create · SigureMo · web-flow · commit 02ca80d307f4 · 2023-12-13T09:03:43.000+08:00
---------

Co-authored-by: SigureMo &lt;sigure.qaq@gmail.com&gt;
diff --git a/test/dygraph_to_static/test_return.py b/test/dygraph_to_static/test_return.py
@@ -17,6 +17,7 @@
 import numpy as np
 from dygraph_to_static_utils import (
     Dy2StTestBase,
+    enable_to_static_guard,
     test_ast_only,
     test_legacy_only,
 )
@@ -281,8 +282,7 @@ def setUp(self):
     def init_dygraph_func(self):
         self.dygraph_func = test_return_base
 
-    def _run(self, to_static=False):
-        paddle.jit.enable_to_static(to_static)
+    def _run(self):
         with base.dygraph.guard():
             res = self.dygraph_func(self.input)
             if isinstance(res, (tuple, list)):
@@ -292,8 +292,9 @@ def _run(self, to_static=False):
             return res
 
     def _test_value_impl(self):
-        dygraph_res = self._run(to_static=False)
-        static_res = self._run(to_static=True)
+        with enable_to_static_guard(False):
+            dygraph_res = self._run()
+        static_res = self._run()
         if isinstance(dygraph_res, tuple):
             self.assertTrue(isinstance(static_res, tuple))
             self.assertEqual(len(dygraph_res), len(static_res))
diff --git a/test/dygraph_to_static/test_save_inference_model.py b/test/dygraph_to_static/test_save_inference_model.py
@@ -230,7 +230,6 @@ def load_and_run_inference(
 class TestPartialProgramRaiseError(Dy2StTestBase):
     @test_ast_only
     def test_param_type(self):
-        paddle.jit.enable_to_static(True)
         x_data = np.random.random((20, 20)).astype('float32')
 
         with base.dygraph.guard(base.CPUPlace()):
diff --git a/test/dygraph_to_static/test_save_load.py b/test/dygraph_to_static/test_save_load.py
@@ -19,6 +19,7 @@
 import numpy as np
 from dygraph_to_static_utils import (
     Dy2StTestBase,
+    enable_to_static_guard,
     test_ast_only,
 )
 from test_fetch_feed import Linear
@@ -72,7 +73,6 @@ def test_save_load_same_result(self):
         x_data = np.random.randn(30, 10, 32).astype('float32')
         batch_num = 3
 
-        paddle.jit.enable_to_static(True)
         x = base.dygraph.to_variable(x_data)
         net = Linear(32, 64)
         adam = Adam(learning_rate=0.1, parameters=net.parameters())
@@ -102,15 +102,15 @@ def test_save_load_same_result(self):
 
         x = base.dygraph.to_variable(x_data)
         # predict output
-        paddle.jit.enable_to_static(False)
-        dygraph_out, dygraph_loss = dygraph_net(x)
-
-        np.testing.assert_allclose(
-            dygraph_out.numpy(), static_out.numpy(), rtol=1e-05
-        )
-        np.testing.assert_allclose(
-            dygraph_loss.numpy(), static_loss.numpy(), rtol=1e-05
-        )
+        with enable_to_static_guard(False):
+            dygraph_out, dygraph_loss = dygraph_net(x)
+
+            np.testing.assert_allclose(
+                dygraph_out.numpy(), static_out.numpy(), rtol=1e-05
+            )
+            np.testing.assert_allclose(
+                dygraph_loss.numpy(), static_loss.numpy(), rtol=1e-05
+            )
 
     def _compute_op_num(self, composite_program):
         if paddle.framework.use_pir_api():
diff --git a/test/dygraph_to_static/test_se_resnet.py b/test/dygraph_to_static/test_se_resnet.py
@@ -22,6 +22,7 @@
 import numpy as np
 from dygraph_to_static_utils import (
     Dy2StTestBase,
+    enable_to_static_guard,
     test_default_and_pir,
 )
 from predictor_utils import PredictorTools
@@ -375,8 +376,6 @@ def tearDown(self):
         self.temp_dir.cleanup()
 
     def train(self, train_reader, to_static):
-        paddle.jit.enable_to_static(to_static)
-
         np.random.seed(SEED)
 
         with base.dygraph.guard(place):
@@ -484,20 +483,22 @@ def train(self, train_reader, to_static):
             )
 
     def predict_dygraph(self, data):
-        paddle.jit.enable_to_static(False)
-        with base.dygraph.guard(place):
-            se_resnext = SeResNeXt()
+        with enable_to_static_guard(False):
+            with base.dygraph.guard(place):
+                se_resnext = SeResNeXt()
 
-            model_dict = paddle.load(self.dy_state_dict_save_path + '.pdparams')
-            se_resnext.set_dict(model_dict)
-            se_resnext.eval()
+                model_dict = paddle.load(
+                    self.dy_state_dict_save_path + '.pdparams'
+                )
+                se_resnext.set_dict(model_dict)
+                se_resnext.eval()
 
-            label = np.random.random([1, 1]).astype("int64")
-            img = base.dygraph.to_variable(data)
-            label = base.dygraph.to_variable(label)
-            pred_res, _, _, _ = se_resnext(img, label)
+                label = np.random.random([1, 1]).astype("int64")
+                img = base.dygraph.to_variable(data)
+                label = base.dygraph.to_variable(label)
+                pred_res, _, _, _ = se_resnext(img, label)
 
-            return pred_res.numpy()
+                return pred_res.numpy()
 
     def predict_static(self, data):
         paddle.enable_static()
@@ -574,9 +575,10 @@ def verify_predict(self):
 
     @test_default_and_pir
     def test_check_result(self):
-        pred_1, loss_1, acc1_1, acc5_1 = self.train(
-            self.train_reader, to_static=False
-        )
+        with enable_to_static_guard(False):
+            pred_1, loss_1, acc1_1, acc5_1 = self.train(
+                self.train_reader, to_static=False
+            )
         pred_2, loss_2, acc1_2, acc5_2 = self.train(
             self.train_reader, to_static=True
         )
diff --git a/test/dygraph_to_static/test_sentiment.py b/test/dygraph_to_static/test_sentiment.py
@@ -15,7 +15,7 @@
 import unittest
 
 import numpy as np
-from dygraph_to_static_utils import Dy2StTestBase
+from dygraph_to_static_utils import Dy2StTestBase, enable_to_static_guard
 from test_lac import DynamicGRU
 
 import paddle
@@ -297,8 +297,7 @@ class Args:
     train_step = 10
 
 
-def train(args, to_static):
-    paddle.jit.enable_to_static(to_static)
+def train(args):
     np.random.seed(SEED)
     paddle.seed(SEED)
     paddle.framework.random._manual_program_seed(SEED)
@@ -373,8 +372,9 @@ def setUp(self):
 
     def train_model(self, model_type='cnn_net'):
         self.args.model_type = model_type
-        st_out = train(self.args, True)
-        dy_out = train(self.args, False)
+        st_out = train(self.args)
+        with enable_to_static_guard(False):
+            dy_out = train(self.args)
         np.testing.assert_allclose(
             dy_out,
             st_out,
diff --git a/test/dygraph_to_static/test_seq2seq.py b/test/dygraph_to_static/test_seq2seq.py
@@ -20,6 +20,7 @@
 import numpy as np
 from dygraph_to_static_utils import (
     Dy2StTestBase,
+    enable_to_static_guard,
     test_legacy_only,
     test_sot_mgs0_only,
 )
@@ -207,14 +208,13 @@ def tearDown(self):
         self.temp_dir.cleanup()
 
     def run_dygraph(self, mode="train", attn_model=False):
-        paddle.jit.enable_to_static(False)
-        if mode == "train":
-            return train(self.args, attn_model)
-        else:
-            return infer(self.args, attn_model)
+        with enable_to_static_guard(False):
+            if mode == "train":
+                return train(self.args, attn_model)
+            else:
+                return infer(self.args, attn_model)
 
     def run_static(self, mode="train", attn_model=False):
-        paddle.jit.enable_to_static(True)
         if mode == "train":
             return train(self.args, attn_model)
         else:
diff --git a/test/dygraph_to_static/test_simnet.py b/test/dygraph_to_static/test_simnet.py
@@ -17,7 +17,11 @@
 import unittest
 
 import numpy as np
-from dygraph_to_static_utils import Dy2StTestBase, test_legacy_and_pt_and_pir
+from dygraph_to_static_utils import (
+    Dy2StTestBase,
+    enable_to_static_guard,
+    test_legacy_and_pt_and_pir,
+)
 from simnet_dygraph_model import BOW, HingeLoss
 
 import paddle
@@ -121,12 +125,10 @@ def __len__(self):
 )
 
 
-def train(conf_dict, to_static):
+def train(conf_dict):
     """
     train process
     """
-    paddle.jit.enable_to_static(to_static)
-
     with unique_name.guard():
         # Get device
         if paddle.is_compiled_with_cuda():
@@ -182,8 +184,10 @@ def test_dygraph_static_same_loss(self):
         if paddle.is_compiled_with_cuda():
             paddle.set_flags({"FLAGS_cudnn_deterministic": True})
         conf_dict = create_conf_dict()
-        dygraph_loss = train(conf_dict, to_static=False)
-        static_loss = train(conf_dict, to_static=True)
+        with enable_to_static_guard(False):
+            dygraph_loss = train(conf_dict)
+
+        static_loss = train(conf_dict)
 
         self.assertEqual(len(dygraph_loss), len(static_loss))
         for i in range(len(dygraph_loss)):
diff --git a/test/dygraph_to_static/test_slice.py b/test/dygraph_to_static/test_slice.py
@@ -19,6 +19,7 @@
 import numpy as np
 from dygraph_to_static_utils import (
     Dy2StTestBase,
+    enable_to_static_guard,
     test_ast_only,
     test_legacy_and_pt_and_pir,
 )
@@ -190,12 +191,17 @@ def tearDown(self):
 
     @test_ast_only
     def test_set_value_with_save(self):
-        paddle.jit.enable_to_static(True)
-        model = paddle.jit.to_static(LayerWithSetValue(input_dim=10, hidden=1))
-        x = paddle.full(shape=[5, 10], fill_value=5.0, dtype="float32")
-        paddle.jit.save(
-            layer=model, path=self.model_path, input_spec=[x], output_spec=None
-        )
+        with enable_to_static_guard(True):
+            model = paddle.jit.to_static(
+                LayerWithSetValue(input_dim=10, hidden=1)
+            )
+            x = paddle.full(shape=[5, 10], fill_value=5.0, dtype="float32")
+            paddle.jit.save(
+                layer=model,
+                path=self.model_path,
+                input_spec=[x],
+                output_spec=None,
+            )
 
 
 class TestSliceSupplementSpecialCase(Dy2StTestBase):
