[Dy2St] enable_to_static_guard 推全 40-45 (#59694)

ooooo-create · DrRyanHuang · SigureMo · web-flow · commit 79e433ab51f9 · 2023-12-13T09:04:19.000+08:00
---------

Co-authored-by: Ryan &lt;44900829+DrRyanHuang@users.noreply.github.com&gt;
Co-authored-by: SigureMo &lt;sigure.qaq@gmail.com&gt;
diff --git a/test/dygraph_to_static/test_tensor_memcpy_on_cpu.py b/test/dygraph_to_static/test_tensor_memcpy_on_cpu.py
@@ -17,6 +17,7 @@
 import numpy as np
 from dygraph_to_static_utils import (
     Dy2StTestBase,
+    enable_to_static_guard,
     test_legacy_and_pt,
     test_legacy_and_pt_and_pir,
 )
@@ -43,19 +44,18 @@ def tensor_copy_to_cuda_with_warning(x, device_id=None, blocking=True):
 
 
 class TestTensorCopyToCpuOnDefaultCPU(Dy2StTestBase):
-    def _run(self, to_static):
-        paddle.jit.enable_to_static(to_static)
+    def _run(self):
         x1 = paddle.ones([1, 2, 3])
         x2 = paddle.jit.to_static(tensor_copy_to_cpu)(x1)
         return x1.place, x2.place, x2.numpy()
 
     @test_legacy_and_pt_and_pir
     def test_tensor_cpu_on_default_cpu(self):
         paddle.framework._set_expected_place(paddle.CPUPlace())
-        dygraph_x1_place, dygraph_place, dygraph_res = self._run(
-            to_static=False
-        )
-        static_x1_place, static_place, static_res = self._run(to_static=True)
+        with enable_to_static_guard(False):
+            dygraph_x1_place, dygraph_place, dygraph_res = self._run()
+
+        static_x1_place, static_place, static_res = self._run()
         np.testing.assert_allclose(dygraph_res, static_res, rtol=1e-05)
         self.assertTrue(dygraph_x1_place.is_cpu_place())
         self.assertTrue(static_x1_place.is_cpu_place())
@@ -64,8 +64,7 @@ def test_tensor_cpu_on_default_cpu(self):
 
 
 class TestTensorCopyToCUDAOnDefaultCPU(Dy2StTestBase):
-    def _run(self, to_static):
-        paddle.jit.enable_to_static(to_static)
+    def _run(self):
         x1 = paddle.ones([1, 2, 3])
         x2 = paddle.jit.to_static(tensor_copy_to_cuda)(x1)
         return x1.place, x2.place, x2.numpy()
@@ -82,10 +81,10 @@ def test_tensor_cuda_on_default_cpu(self):
         See ConstructDeviceContext() in interpreter_util.cc.
         """
         paddle.framework._set_expected_place(paddle.CPUPlace())
-        dygraph_x1_place, dygraph_place, dygraph_res = self._run(
-            to_static=False
-        )
-        static_x1_place, static_place, static_res = self._run(to_static=True)
+        with enable_to_static_guard(False):
+            dygraph_x1_place, dygraph_place, dygraph_res = self._run()
+
+        static_x1_place, static_place, static_res = self._run()
         np.testing.assert_allclose(dygraph_res, static_res, rtol=1e-05)
         self.assertTrue(dygraph_x1_place.is_cpu_place())
         self.assertTrue(static_x1_place.is_cpu_place())
@@ -94,8 +93,7 @@ def test_tensor_cuda_on_default_cpu(self):
 
 
 class TestTensorCopyToCUDAWithWarningOnCPU(unittest.TestCase):
-    def _run(self, to_static):
-        paddle.jit.enable_to_static(to_static)
+    def _run(self):
         x1 = paddle.ones([1, 2, 3])
         x2 = paddle.jit.to_static(tensor_copy_to_cuda_with_warning)(
             x1, device_id=1, blocking=False
diff --git a/test/dygraph_to_static/test_tensor_memcpy_on_gpu.py b/test/dygraph_to_static/test_tensor_memcpy_on_gpu.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,
 )
 
@@ -43,8 +44,7 @@ def tensor_copy_to_cuda_with_warning(x, device_id=None, blocking=True):
 
 
 class TestTensorCopyToCpuOnDefaultGPU(Dy2StTestBase):
-    def _run(self, to_static):
-        paddle.jit.enable_to_static(to_static)
+    def _run(self):
         x1 = paddle.ones([1, 2, 3])
         x2 = paddle.jit.to_static(tensor_copy_to_cpu)(x1)
         return x1.place, x2.place, x2.numpy()
@@ -55,10 +55,10 @@ def test_tensor_cpu_on_default_gpu(self):
             return
         place = paddle.CUDAPlace(int(os.environ.get('FLAGS_selected_gpus', 0)))
         paddle.framework._set_expected_place(place)
-        dygraph_x1_place, dygraph_place, dygraph_res = self._run(
-            to_static=False
-        )
-        static_x1_place, static_place, static_res = self._run(to_static=True)
+        with enable_to_static_guard(False):
+            dygraph_x1_place, dygraph_place, dygraph_res = self._run()
+
+        static_x1_place, static_place, static_res = self._run()
         np.testing.assert_allclose(dygraph_res, static_res, rtol=1e-05)
         self.assertTrue(dygraph_x1_place.is_gpu_place())
         self.assertTrue(static_x1_place.is_gpu_place())
@@ -67,8 +67,7 @@ def test_tensor_cpu_on_default_gpu(self):
 
 
 class TestTensorCopyToCUDAOnDefaultGPU(Dy2StTestBase):
-    def _run(self, to_static):
-        paddle.jit.enable_to_static(to_static)
+    def _run(self):
         x1 = paddle.ones([1, 2, 3])
         x2 = paddle.jit.to_static(tensor_copy_to_cuda)(x1)
         return x1.place, x2.place, x2.numpy()
@@ -79,10 +78,10 @@ def test_tensor_cuda_on_default_gpu(self):
             return
         place = paddle.CUDAPlace(int(os.environ.get('FLAGS_selected_gpus', 0)))
         paddle.framework._set_expected_place(place)
-        dygraph_x1_place, dygraph_place, dygraph_res = self._run(
-            to_static=False
-        )
-        static_x1_place, static_place, static_res = self._run(to_static=True)
+        with enable_to_static_guard(False):
+            dygraph_x1_place, dygraph_place, dygraph_res = self._run()
+
+        static_x1_place, static_place, static_res = self._run()
         np.testing.assert_allclose(dygraph_res, static_res, rtol=1e-05)
         self.assertTrue(dygraph_x1_place.is_gpu_place())
         self.assertTrue(static_x1_place.is_gpu_place())
@@ -91,8 +90,7 @@ def test_tensor_cuda_on_default_gpu(self):
 
 
 class TestTensorCopyToCUDAWithWarningOnGPU(unittest.TestCase):
-    def _run(self, to_static):
-        paddle.jit.enable_to_static(to_static)
+    def _run(self):
         x1 = paddle.ones([1, 2, 3])
         x2 = paddle.jit.to_static(tensor_copy_to_cuda_with_warning)(
             x1, device_id=1, blocking=False
diff --git a/test/dygraph_to_static/test_tensor_methods.py b/test/dygraph_to_static/test_tensor_methods.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_and_pt_and_pir,
 )
@@ -31,16 +32,17 @@ def tensor_clone(x):
 
 
 class TestTensorClone(Dy2StTestBase):
-    def _run(self, to_static):
-        paddle.jit.enable_to_static(to_static)
+    def _run(self):
         x = paddle.ones([1, 2, 3])
         return paddle.jit.to_static(tensor_clone)(x).numpy()
 
     @test_legacy_and_pt_and_pir
     def test_tensor_clone(self):
         paddle.disable_static()
-        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()
         np.testing.assert_allclose(dygraph_res, static_res, rtol=1e-05)
 
 
@@ -51,8 +53,7 @@ def tensor_numpy(x):
 
 
 class TestTensorDygraphOnlyMethodError(Dy2StTestBase):
-    def _run(self, to_static):
-        paddle.jit.enable_to_static(to_static)
+    def _run(self):
         x = paddle.zeros([2, 2])
         y = paddle.jit.to_static(tensor_numpy)(x)
         return y.numpy()
@@ -61,9 +62,11 @@ def _run(self, to_static):
     @test_legacy_and_pt_and_pir
     def test_to_static_numpy_report_error(self):
         paddle.disable_static()
-        dygraph_res = self._run(to_static=False)
+        with enable_to_static_guard(False):
+            dygraph_res = self._run()
+
         with self.assertRaises(AssertionError):
-            static_res = self._run(to_static=True)
+            static_res = self._run()
 
 
 def tensor_item(x):
@@ -73,16 +76,17 @@ def tensor_item(x):
 
 
 class TestTensorItem(Dy2StTestBase):
-    def _run(self, to_static):
-        paddle.jit.enable_to_static(to_static)
+    def _run(self):
         x = paddle.ones([1])
         return paddle.jit.to_static(tensor_item)(x)
 
     @test_legacy_and_pt_and_pir
     def test_tensor_clone(self):
         paddle.disable_static()
-        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()
         np.testing.assert_allclose(dygraph_res, static_res)
 
 
@@ -95,7 +99,6 @@ def tensor_size(x):
 
 class TestTensorSize(Dy2StTestBase):
     def _run(self, to_static):
-        paddle.jit.enable_to_static(to_static)
         x = paddle.ones([1, 2, 3])
         if not to_static:
             return tensor_size(x)
@@ -107,7 +110,9 @@ def _run(self, to_static):
     @test_legacy_and_pt_and_pir
     def test_tensor_clone(self):
         paddle.disable_static()
-        dygraph_res = self._run(to_static=False)
+        with enable_to_static_guard(False):
+            dygraph_res = self._run(to_static=False)
+
         static_res = self._run(to_static=True)
         np.testing.assert_allclose(dygraph_res, static_res, rtol=1e-5)
 
@@ -118,17 +123,17 @@ def true_div(x, y):
 
 
 class TestTrueDiv(Dy2StTestBase):
-    def _run(self, to_static):
-        paddle.jit.enable_to_static(to_static)
+    def _run(self):
         x = paddle.to_tensor([3], dtype='int64')
         y = paddle.to_tensor([4], dtype='int64')
         return paddle.jit.to_static(true_div)(x, y).numpy()
 
     @test_legacy_and_pt_and_pir
-    def test_ture_div(self):
+    def test_true_div(self):
         paddle.disable_static()
-        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()
         np.testing.assert_allclose(dygraph_res, static_res, rtol=1e-5)
 
 
diff --git a/test/dygraph_to_static/test_tsm.py b/test/dygraph_to_static/test_tsm.py
@@ -21,6 +21,7 @@
 import numpy as np
 from dygraph_to_static_utils import (
     Dy2StTestBase,
+    enable_to_static_guard,
     test_default_and_pir,
 )
 from tsm_config_utils import merge_configs, parse_config, print_configs
@@ -292,9 +293,7 @@ def create_optimizer(cfg, params):
     return optimizer
 
 
-def train(args, fake_data_reader, to_static):
-    paddle.jit.enable_to_static(to_static)
-
+def train(args, fake_data_reader):
     config = parse_config(args.config)
     train_config = merge_configs(config, 'train', vars(args))
     valid_config = merge_configs(config, 'valid', vars(args))
@@ -381,8 +380,10 @@ def test_dygraph_static_same_loss(self):
             paddle.set_flags({"FLAGS_cudnn_deterministic": True})
         args = parse_args()
         fake_data_reader = FakeDataReader("train", parse_config(args.config))
-        dygraph_loss = train(args, fake_data_reader, to_static=False)
-        static_loss = train(args, fake_data_reader, to_static=True)
+        with enable_to_static_guard(False):
+            dygraph_loss = train(args, fake_data_reader)
+
+        static_loss = train(args, fake_data_reader)
         np.testing.assert_allclose(dygraph_loss, static_loss, rtol=1e-05)
 
 
diff --git a/test/dygraph_to_static/test_word2vec.py b/test/dygraph_to_static/test_word2vec.py
@@ -19,6 +19,7 @@
 import numpy as np
 from dygraph_to_static_utils import (
     Dy2StTestBase,
+    enable_to_static_guard,
     test_legacy_and_pt_and_pir,
 )
 
@@ -275,9 +276,7 @@ def forward(self, center_words, target_words, label):
 total_steps = len(dataset) * epoch_num // batch_size
 
 
-def train(to_static):
-    paddle.jit.enable_to_static(to_static)
-
+def train():
     random.seed(0)
     np.random.seed(0)
 
@@ -322,8 +321,10 @@ def train(to_static):
 class TestWord2Vec(Dy2StTestBase):
     @test_legacy_and_pt_and_pir
     def test_dygraph_static_same_loss(self):
-        dygraph_loss = train(to_static=False)
-        static_loss = train(to_static=True)
+        with enable_to_static_guard(False):
+            dygraph_loss = train()
+
+        static_loss = train()
         np.testing.assert_allclose(dygraph_loss, static_loss, rtol=1e-05)
 
 
diff --git a/test/dygraph_to_static/test_yolov3.py b/test/dygraph_to_static/test_yolov3.py
@@ -19,6 +19,7 @@
 import numpy as np
 from dygraph_to_static_utils import (
     Dy2StTestBase,
+    enable_to_static_guard,
     test_default_and_pir,
 )
 from yolov3 import YOLOv3, cfg
@@ -81,9 +82,7 @@ def generator():
 fake_data_reader = FakeDataReader()
 
 
-def train(to_static):
-    paddle.jit.enable_to_static(to_static)
-
+def train():
     random.seed(0)
     np.random.seed(0)
 
@@ -172,8 +171,10 @@ def train(to_static):
 class TestYolov3(Dy2StTestBase):
     @test_default_and_pir
     def test_dygraph_static_same_loss(self):
-        dygraph_loss = train(to_static=False)
-        static_loss = train(to_static=True)
+        with enable_to_static_guard(False):
+            dygraph_loss = train()
+
+        static_loss = train()
         np.testing.assert_allclose(
             dygraph_loss, static_loss, rtol=0.001, atol=1e-05
         )
