[Dy2St][tests][17-22] paddle.jit.enable_to_static->enable_to_static_guard (#59717)

gouzil · web-flow · commit 77227c485e32 · 2023-12-13T21:04:03.000+08:00
diff --git a/test/dygraph_to_static/test_logical.py b/test/dygraph_to_static/test_logical.py
@@ -18,10 +18,9 @@
 import unittest
 
 import numpy as np
-from dygraph_to_static_utils import Dy2StTestBase
+from dygraph_to_static_utils import Dy2StTestBase, enable_to_static_guard
 
 import paddle
-from paddle import base
 from paddle.jit.dy2static.logical_transformer import cmpop_node_to_str
 from paddle.utils import gast
 
@@ -162,21 +161,15 @@ def test_shape_not_equal(x):
 class TestLogicalBase(Dy2StTestBase):
     def setUp(self):
         self.input = np.array([3]).astype('int32')
-        self.place = (
-            paddle.CUDAPlace(0)
-            if base.is_compiled_with_cuda()
-            else paddle.CPUPlace()
-        )
         self._set_test_func()
 
     def _set_test_func(self):
         raise NotImplementedError(
             "Method 'set_test_func' should be implemented."
         )
 
-    def _run(self, to_static):
-        paddle.jit.enable_to_static(to_static)
-        with base.dygraph.guard(self.place):
+    def _run(self, to_static: bool):
+        with enable_to_static_guard(to_static):
             result = paddle.jit.to_static(self.dygraph_func)(self.input)
             return result.numpy()
 
diff --git a/test/dygraph_to_static/test_lstm.py b/test/dygraph_to_static/test_lstm.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,
 )
@@ -58,23 +59,21 @@ def tearDown(self):
         self.temp_dir.cleanup()
 
     def run_lstm(self, to_static):
-        paddle.jit.enable_to_static(to_static)
+        with enable_to_static_guard(to_static):
+            paddle.static.default_main_program().random_seed = 1001
+            paddle.static.default_startup_program().random_seed = 1001
 
-        paddle.static.default_main_program().random_seed = 1001
-        paddle.static.default_startup_program().random_seed = 1001
-
-        net = paddle.jit.to_static(Net(12, 2))
-        x = paddle.zeros((2, 10, 12))
-        y = net(x)
-        return y.numpy()
+            net = paddle.jit.to_static(Net(12, 2))
+            x = paddle.zeros((2, 10, 12))
+            y = net(x)
+            return y.numpy()
 
     def test_lstm_to_static(self):
         dygraph_out = self.run_lstm(to_static=False)
         static_out = self.run_lstm(to_static=True)
         np.testing.assert_allclose(dygraph_out, static_out, rtol=1e-05)
 
     def save_in_eval(self, with_training: bool):
-        paddle.jit.enable_to_static(True)
         net = Net(12, 2)
         x = paddle.randn((2, 10, 12))
         if with_training:
@@ -150,7 +149,6 @@ def tearDown(self):
 
     @test_legacy_and_pt_and_pir
     def test_save_in_eval(self):
-        paddle.jit.enable_to_static(True)
         net = paddle.jit.to_static(LinearNet())
         x = paddle.randn((2, 10))
         x.stop_gradient = False
diff --git a/test/dygraph_to_static/test_mobile_net.py b/test/dygraph_to_static/test_mobile_net.py
@@ -21,6 +21,7 @@
 import numpy as np
 from dygraph_to_static_utils import (
     Dy2StTestBase,
+    enable_to_static_guard,
     test_legacy_and_pt_and_pir,
 )
 from predictor_utils import PredictorTools
@@ -509,8 +510,6 @@ class Args:
 
 
 def train_mobilenet(args, to_static):
-    paddle.jit.enable_to_static(to_static)
-
     with unique_name.guard():
         np.random.seed(SEED)
         paddle.seed(SEED)
@@ -632,23 +631,23 @@ def predict_static(args, data):
 
 
 def predict_dygraph(args, data):
-    paddle.jit.enable_to_static(False)
-    if args.model == "MobileNetV1":
-        model = paddle.jit.to_static(
-            MobileNetV1(class_dim=args.class_dim, scale=1.0)
-        )
-    elif args.model == "MobileNetV2":
-        model = paddle.jit.to_static(
-            MobileNetV2(class_dim=args.class_dim, scale=1.0)
-        )
-    # load dygraph trained parameters
-    model_dict = paddle.load(args.dy_state_dict_save_path + '.pdparams')
-    model.set_dict(model_dict)
-    model.eval()
+    with enable_to_static_guard(False):
+        if args.model == "MobileNetV1":
+            model = paddle.jit.to_static(
+                MobileNetV1(class_dim=args.class_dim, scale=1.0)
+            )
+        elif args.model == "MobileNetV2":
+            model = paddle.jit.to_static(
+                MobileNetV2(class_dim=args.class_dim, scale=1.0)
+            )
+        # load dygraph trained parameters
+        model_dict = paddle.load(args.dy_state_dict_save_path + '.pdparams')
+        model.set_dict(model_dict)
+        model.eval()
 
-    pred_res = model(base.dygraph.to_variable(data))
+        pred_res = model(base.dygraph.to_variable(data))
 
-    return pred_res.numpy()
+        return pred_res.numpy()
 
 
 def predict_dygraph_jit(args, data):
@@ -689,7 +688,8 @@ def train(self, model_name, to_static):
         self.args.dy_state_dict_save_path = os.path.join(
             self.temp_dir.name, model_name + ".dygraph"
         )
-        out = train_mobilenet(self.args, to_static)
+        with enable_to_static_guard(to_static):
+            out = train_mobilenet(self.args, to_static)
         return out
 
     def assert_same_loss(self, model_name):
diff --git a/test/dygraph_to_static/test_param_guard.py b/test/dygraph_to_static/test_param_guard.py
@@ -18,6 +18,7 @@
 import numpy as np
 from dygraph_to_static_utils import (
     Dy2StTestBase,
+    enable_to_static_guard,
     test_legacy_and_pt_and_pir,
 )
 
@@ -60,25 +61,25 @@ def setUp(self):
         self.seed = 2021
         self.iter_num = 5
 
-    def train(self, is_iter, to_static):
+    def train(self, is_iter, to_static: bool):
         paddle.seed(self.seed)
         np.random.seed(self.seed)
-        paddle.jit.enable_to_static(to_static)
-        if is_iter:
-            net = paddle.jit.to_static(NetWithParameterList(10, 3))
-        else:
-            net = paddle.jit.to_static(NetWithParameterListIter(10, 3))
-        sgd = paddle.optimizer.SGD(0.1, parameters=net.parameters())
-
-        for batch_id in range(self.iter_num):
-            x = paddle.rand([4, 10], dtype='float32')
-            out = net(x)
-            loss = paddle.mean(out)
-            loss.backward()
-            sgd.step()
-            sgd.clear_grad()
-
-        return loss
+        with enable_to_static_guard(to_static):
+            if is_iter:
+                net = paddle.jit.to_static(NetWithParameterList(10, 3))
+            else:
+                net = paddle.jit.to_static(NetWithParameterListIter(10, 3))
+            sgd = paddle.optimizer.SGD(0.1, parameters=net.parameters())
+
+            for batch_id in range(self.iter_num):
+                x = paddle.rand([4, 10], dtype='float32')
+                out = net(x)
+                loss = paddle.mean(out)
+                loss.backward()
+                sgd.step()
+                sgd.clear_grad()
+
+            return loss
 
     @test_legacy_and_pt_and_pir
     def test_parameter_list(self):
@@ -114,23 +115,23 @@ def setUp(self):
     def init_net(self):
         self.net = paddle.jit.to_static(NetWithRawParamList(10, 3))
 
-    def train(self, to_static):
+    def train(self, to_static: bool):
         paddle.seed(self.seed)
         np.random.seed(self.seed)
-        paddle.jit.enable_to_static(to_static)
-        self.init_net()
+        with enable_to_static_guard(to_static):
+            self.init_net()
 
-        sgd = paddle.optimizer.SGD(0.1, parameters=self.net.parameters())
+            sgd = paddle.optimizer.SGD(0.1, parameters=self.net.parameters())
 
-        for batch_id in range(self.iter_num):
-            x = paddle.rand([4, 10], dtype='float32')
-            out = self.net(x)
-            loss = paddle.mean(out)
-            loss.backward()
-            sgd.step()
-            sgd.clear_grad()
+            for batch_id in range(self.iter_num):
+                x = paddle.rand([4, 10], dtype='float32')
+                out = self.net(x)
+                loss = paddle.mean(out)
+                loss.backward()
+                sgd.step()
+                sgd.clear_grad()
 
-        return loss
+            return loss
 
     @test_legacy_and_pt_and_pir
     def test_parameter_list(self):
diff --git a/test/dygraph_to_static/test_print.py b/test/dygraph_to_static/test_print.py
@@ -17,6 +17,7 @@
 import numpy
 from dygraph_to_static_utils import (
     Dy2StTestBase,
+    enable_to_static_guard,
     test_legacy_and_pt_and_pir,
 )
 
@@ -80,20 +81,14 @@ def dyfunc_print_with_kwargs(x):
 class TestPrintBase(Dy2StTestBase):
     def setUp(self):
         self.input = numpy.ones(5).astype("int32")
-        self.place = (
-            paddle.CUDAPlace(0)
-            if paddle.is_compiled_with_cuda()
-            else paddle.CPUPlace()
-        )
         self.set_test_func()
 
     def set_test_func(self):
         raise NotImplementedError("Print test should implement set_test_func")
 
-    def _run(self, to_static):
-        paddle.jit.enable_to_static(to_static)
-
-        paddle.jit.to_static(self.dygraph_func)(self.input)
+    def _run(self, to_static: bool):
+        with enable_to_static_guard(to_static):
+            paddle.jit.to_static(self.dygraph_func)(self.input)
 
     def get_dygraph_output(self):
         self._run(to_static=False)
diff --git a/test/dygraph_to_static/test_program_translator.py b/test/dygraph_to_static/test_program_translator.py
@@ -35,7 +35,6 @@
 
 import paddle
 import paddle.jit.dy2static as _jst
-from paddle.jit.api import to_static
 from paddle.jit.dy2static.utils import func_to_source_code
 from paddle.utils import gast
 
@@ -228,16 +227,18 @@ def setUp(self):
     @test_ast_only
     @test_legacy_and_pt_and_pir
     def test_raise_error(self):
-        net = to_static(full_graph=True)(NetWithError())
+        net = paddle.jit.to_static(full_graph=True)(NetWithError())
         with self.assertRaises(ValueError):
             net(paddle.to_tensor(self.x))
 
     @test_legacy_and_pt_and_pir
     def test_enable_disable_to_static(self):
-        static_output = to_static(decorated_simple_func)(self.x, self.weight)
+        static_output = paddle.jit.to_static(decorated_simple_func)(
+            self.x, self.weight
+        )
 
         with enable_to_static_guard(False):
-            dygraph_output = to_static(decorated_simple_func)(
+            dygraph_output = paddle.jit.to_static(decorated_simple_func)(
                 self.x, self.weight
             )
             np.testing.assert_allclose(
