[Dy2St][tests][23-29] paddle.jit.enable_to_static->enable_to_static_guard (#59729)

gouzil · web-flow · commit bd0039a8e080 · 2023-12-13T21:05:08.000+08:00
diff --git a/test/dygraph_to_static/test_ptb_lm.py b/test/dygraph_to_static/test_ptb_lm.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,
+)
 
 import paddle
 from paddle.base.framework import unique_name
@@ -310,20 +314,19 @@ def train():
 
 
 def train_dygraph():
-    paddle.jit.enable_to_static(False)
-    return train()
+    with enable_to_static_guard(False):
+        return train()
 
 
 def train_static():
-    paddle.jit.enable_to_static(True)
     return train()
 
 
 class TestPtb(Dy2StTestBase):
     @test_legacy_and_pt_and_pir
     def test_check_result(self):
-        loss_1, hidden_1, cell_1 = train_static()
-        loss_2, hidden_2, cell_2 = train_dygraph()
+        loss_1, hidden_1, cell_1 = train_dygraph()
+        loss_2, hidden_2, cell_2 = train_static()
 
         np.testing.assert_allclose(loss_1, loss_2, rtol=1e-05)
         np.testing.assert_allclose(hidden_1, hidden_2, rtol=1e-05)
diff --git a/test/dygraph_to_static/test_pylayer.py b/test/dygraph_to_static/test_pylayer.py
@@ -17,6 +17,8 @@
 import sys
 from pathlib import Path
 
+from dygraph_to_static_utils import enable_to_static_guard
+
 sys.path.append(
     str(Path(__file__).absolute().parent.parent.joinpath("legacy_test"))
 )
@@ -265,18 +267,18 @@ def forward(self, x):
 class TestPyLayerBase(unittest.TestCase):
     def setUp(self):
         self.place = "gpu" if paddle.is_compiled_with_cuda() else "cpu"
-        self.to_static = False
+        self.to_static: bool = False
 
     def _run(self, *input_args, **input_kwargs):
         assert getattr(
             self, "dygraph_func", None
         ), "Please setting `self.dygraph_func` before calling `self._run`"
 
-        paddle.jit.enable_to_static(self.to_static)
-        paddle.set_device(self.place)
-        result = self.dygraph_func(*input_args, **input_kwargs)
-        result.mean().backward()
-        return result
+        with enable_to_static_guard(self.to_static):
+            paddle.set_device(self.place)
+            result = self.dygraph_func(*input_args, **input_kwargs)
+            result.mean().backward()
+            return result
 
     def _run_dygraph(self, *args, **kwargs):
         self.to_static = False
@@ -516,12 +518,10 @@ class PyLayerTrainHelper(unittest.TestCase):
     def setUp(self):
         self.place = "gpu" if paddle.is_compiled_with_cuda() else "cpu"
 
-    def _run_train(self, to_static, layer_builder, build_strategy=None):
+    def _run_train(self, to_static: bool, layer_builder, build_strategy=None):
         """
         Tests model decorated by `dygraph_to_static_output` in static graph mode. For users, the model is defined in dygraph mode and trained in static graph mode.
         """
-        paddle.jit.enable_to_static(to_static)
-
         paddle.set_device(self.place)
         np.random.seed(SEED)
         paddle.seed(SEED)
diff --git a/test/dygraph_to_static/test_reinforcement_learning.py b/test/dygraph_to_static/test_reinforcement_learning.py
@@ -20,12 +20,12 @@
 import numpy as np
 from dygraph_to_static_utils import (
     Dy2StTestBase,
+    enable_to_static_guard,
     test_legacy_and_pt_and_pir,
 )
 
 import paddle
 import paddle.nn.functional as F
-from paddle import base
 from paddle.base.dygraph import to_variable
 from paddle.nn import Layer
 
@@ -61,13 +61,11 @@ class Args:
     train_step = 10
 
 
-def train(args, place, to_static):
-    paddle.jit.enable_to_static(to_static)
+def train(args, to_static: bool):
+    with enable_to_static_guard(to_static):
+        env = gym.make('CartPole-v0')
+        env.reset(seed=SEED)
 
-    env = gym.make('CartPole-v0')
-    env.reset(seed=SEED)
-
-    with base.dygraph.guard(place):
         paddle.seed(SEED)
         paddle.framework.random._manual_program_seed(SEED)
         local_random = np.random.RandomState(SEED)
@@ -206,17 +204,12 @@ def finish_episode():
 
 class TestDeclarative(Dy2StTestBase):
     def setUp(self):
-        self.place = (
-            paddle.CUDAPlace(0)
-            if paddle.is_compiled_with_cuda()
-            else paddle.CPUPlace()
-        )
         self.args = Args()
 
     @test_legacy_and_pt_and_pir
     def test_train(self):
-        st_out = train(self.args, self.place, to_static=True)
-        dy_out = train(self.args, self.place, to_static=False)
+        st_out = train(self.args, to_static=True)
+        dy_out = train(self.args, to_static=False)
         np.testing.assert_allclose(st_out, dy_out, rtol=1e-05)
 
 
diff --git a/test/dygraph_to_static/test_resnet_amp.py b/test/dygraph_to_static/test_resnet_amp.py
@@ -16,34 +16,37 @@
 import unittest
 
 import numpy as np
-from dygraph_to_static_utils import Dy2StTestBase, test_default_mode_only
+from dygraph_to_static_utils import (
+    Dy2StTestBase,
+    enable_to_static_guard,
+    test_default_mode_only,
+)
 from test_resnet import SEED, ResNet, optimizer_setting
 
 import paddle
-from paddle import base
 from paddle.base import core
 
 # NOTE: Reduce batch_size from 8 to 2 to avoid unittest timeout.
 batch_size = 2
 epoch_num = 1
-place = base.CUDAPlace(0) if base.is_compiled_with_cuda() else base.CPUPlace()
+place = (
+    paddle.CUDAPlace(0) if paddle.is_compiled_with_cuda() else paddle.CPUPlace()
+)
 
 
-if base.is_compiled_with_cuda():
-    base.set_flags({'FLAGS_cudnn_deterministic': True})
+if paddle.is_compiled_with_cuda():
+    paddle.set_flags({'FLAGS_cudnn_deterministic': True})
 
 
-def train(to_static, build_strategy=None):
+def train(build_strategy=None):
     """
     Tests model decorated by `dygraph_to_static_output` in static graph mode. For users, the model is defined in dygraph mode and trained in static graph mode.
     """
     np.random.seed(SEED)
     paddle.seed(SEED)
     paddle.framework.random._manual_program_seed(SEED)
 
-    resnet = ResNet()
-    if to_static:
-        resnet = paddle.jit.to_static(resnet, build_strategy=build_strategy)
+    resnet = paddle.jit.to_static(ResNet(), build_strategy=build_strategy)
     optimizer = optimizer_setting(parameter_list=resnet.parameters())
     scaler = paddle.amp.GradScaler(init_loss_scaling=1024)
 
@@ -109,9 +112,9 @@ def train(to_static, build_strategy=None):
 
 
 class TestResnet(Dy2StTestBase):
-    def train(self, to_static):
-        paddle.jit.enable_to_static(to_static)
-        return train(to_static)
+    def train(self, to_static: bool):
+        with enable_to_static_guard(to_static):
+            return train()
 
     @test_default_mode_only
     def test_resnet(self):
diff --git a/test/dygraph_to_static/test_resnet_pure_fp16.py b/test/dygraph_to_static/test_resnet_pure_fp16.py
@@ -16,7 +16,11 @@
 import unittest
 
 import numpy as np
-from dygraph_to_static_utils import Dy2StTestBase, test_default_mode_only
+from dygraph_to_static_utils import (
+    Dy2StTestBase,
+    enable_to_static_guard,
+    test_default_mode_only,
+)
 from test_resnet import SEED, ResNet, optimizer_setting
 
 import paddle
@@ -28,21 +32,19 @@
 epoch_num = 1
 
 
-if base.is_compiled_with_cuda():
-    base.set_flags({'FLAGS_cudnn_deterministic': True})
+if paddle.is_compiled_with_cuda():
+    paddle.set_flags({'FLAGS_cudnn_deterministic': True})
 
 
-def train(to_static, build_strategy=None):
+def train(build_strategy=None):
     """
     Tests model decorated by `dygraph_to_static_output` in static graph mode. For users, the model is defined in dygraph mode and trained in static graph mode.
     """
     np.random.seed(SEED)
     paddle.seed(SEED)
     paddle.framework.random._manual_program_seed(SEED)
 
-    resnet = ResNet()
-    if to_static:
-        resnet = paddle.jit.to_static(resnet, build_strategy=build_strategy)
+    resnet = paddle.jit.to_static(ResNet(), build_strategy=build_strategy)
     optimizer = optimizer_setting(parameter_list=resnet.parameters())
     scaler = paddle.amp.GradScaler(init_loss_scaling=1024)
 
@@ -113,13 +115,13 @@ def train(to_static, build_strategy=None):
 
 
 class TestResnet(Dy2StTestBase):
-    def train(self, to_static):
-        paddle.jit.enable_to_static(to_static)
-        build_strategy = paddle.static.BuildStrategy()
-        # Why set `build_strategy.enable_inplace = False` here?
-        # Because we find that this PASS strategy of PE makes dy2st training loss unstable.
-        build_strategy.enable_inplace = False
-        return train(to_static, build_strategy)
+    def train(self, to_static: bool):
+        with enable_to_static_guard(to_static):
+            build_strategy = paddle.static.BuildStrategy()
+            # Why set `build_strategy.enable_inplace = False` here?
+            # Because we find that this PASS strategy of PE makes dy2st training loss unstable.
+            build_strategy.enable_inplace = False
+            return train(build_strategy)
 
     @test_default_mode_only
     def test_resnet(self):
