[2.6][Dy2St] Decrease test_train_step time and disable some uts on CPU(#59867) (#59923)

SigureMo · web-flow · commit d47d31e82dc2 · 2023-12-12T17:11:37.000+08:00
* [Dy2St] Decrease `test_train_step` time (#59867) * [Dy2St] Decrease `test_train_step` time * rename back * [Dy2St] Disable `test_resnet` and `test_build_strategy` on CPU tests (#59742) * [Dy2St] Decrease `test_resnet` time * temp rename to resnetx * increase timeout of test_resnet * remove pir test case * dec timeout * inc timeout * disable on CPU * rename to merge * rename to resnetx * rename to resnet * fix import * fix v2 * fix test_build_strategy * rename to test_build_strategyx * rename back * only,cherry-pick; test=document_fix
diff --git a/test/dygraph_to_static/CMakeLists.txt b/test/dygraph_to_static/CMakeLists.txt
@@ -25,6 +25,9 @@ if(NOT WITH_GPU)
   # We should remove this after fix the performance issue.
   list(REMOVE_ITEM TEST_OPS test_train_step_resnet18_adam)
   list(REMOVE_ITEM TEST_OPS test_train_step_resnet18_sgd)
+  # disable some model test on CPU to avoid timeout
+  list(REMOVE_ITEM TEST_OPS test_resnet)
+  list(REMOVE_ITEM TEST_OPS test_build_strategy)
 endif()
 
 foreach(TEST_OP ${TEST_OPS})
@@ -43,19 +46,14 @@ set_tests_properties(test_reinforcement_learning PROPERTIES TIMEOUT 120)
 set_tests_properties(test_transformer PROPERTIES TIMEOUT 200)
 set_tests_properties(test_bmn PROPERTIES TIMEOUT 300)
 set_tests_properties(test_bert PROPERTIES TIMEOUT 240)
-#set_tests_properties(test_mnist PROPERTIES TIMEOUT 120)
-set_tests_properties(test_build_strategy PROPERTIES TIMEOUT 120)
 
 if(NOT WIN32)
   set_tests_properties(test_tsm PROPERTIES TIMEOUT 900)
-  set_tests_properties(test_resnet PROPERTIES TIMEOUT 300)
 endif()
 
 if(APPLE)
   set_tests_properties(test_bmn PROPERTIES TIMEOUT 300)
-  set_tests_properties(test_build_strategy PROPERTIES TIMEOUT 300)
   set_tests_properties(test_mobile_net PROPERTIES TIMEOUT 300)
-  set_tests_properties(test_resnet PROPERTIES TIMEOUT 300)
 endif()
 
 if(WITH_GPU)
diff --git a/test/dygraph_to_static/dygraph_to_static_utils.py b/test/dygraph_to_static/dygraph_to_static_utils.py
@@ -381,23 +381,23 @@ def test_legacy_and_pt_and_pir(fn):
     return fn
 
 
+# Some decorators for save CI time
 def test_default_mode_only(fn):
     # Some unittests has high time complexity, we only test them with default mode
     fn = set_to_static_mode(ToStaticMode.SOT)(fn)
     fn = set_ir_mode(IrMode.PT)(fn)
     return fn
 
 
-def test_sot_with_pir_only(fn):
+def test_default_and_pir(fn):
+    # Some unittests has high time complexity, we only test them with default mode
     fn = set_to_static_mode(ToStaticMode.SOT)(fn)
-    fn = set_ir_mode(IrMode.PIR)(fn)
+    fn = set_ir_mode(IrMode.PT | IrMode.PIR)(fn)
     return fn
 
 
-def test_default_and_pir(fn):
-    # Some unittests has high time complexity, we only test them with default mode
+def test_sot_mgs0_only(fn):
     fn = set_to_static_mode(ToStaticMode.SOT)(fn)
-    fn = set_ir_mode(IrMode.PT | IrMode.PIR)(fn)
     return fn
 
 
diff --git a/test/dygraph_to_static/test_build_strategy.py b/test/dygraph_to_static/test_build_strategy.py
@@ -18,9 +18,8 @@
 from dygraph_to_static_utils import (
     Dy2StTestBase,
     enable_to_static_guard,
-    test_ast_only,
-    test_default_and_pir,
-    test_pt_only,
+    test_default_mode_only,
+    test_legacy_and_pt_and_pir,
 )
 from test_resnet import ResNetHelper
 
@@ -36,7 +35,7 @@ def setUp(self):
         self.build_strategy.enable_addto = True
         self.resnet_helper = ResNetHelper()
         # NOTE: for enable_addto
-        paddle.base.set_flags({"FLAGS_max_inplace_grad_add": 8})
+        paddle.set_flags({"FLAGS_max_inplace_grad_add": 8})
 
     def train(self, to_static):
         with enable_to_static_guard(to_static):
@@ -67,8 +66,7 @@ def verify_predict(self):
             err_msg=f'predictor_pre:\n {predictor_pre}\n, st_pre: \n{st_pre}.',
         )
 
-    @test_ast_only
-    @test_pt_only
+    @test_default_mode_only
     def test_resnet(self):
         static_loss = self.train(to_static=True)
         dygraph_loss = self.train(to_static=False)
@@ -80,19 +78,18 @@ def test_resnet(self):
         )
         self.verify_predict()
 
-    @test_ast_only
-    @test_pt_only
+    @test_default_mode_only
     def test_in_static_mode_mkldnn(self):
-        paddle.base.set_flags({'FLAGS_use_mkldnn': True})
+        paddle.set_flags({'FLAGS_use_mkldnn': True})
         try:
             if paddle.base.core.is_compiled_with_mkldnn():
                 self.resnet_helper.train(True, self.build_strategy)
         finally:
-            paddle.base.set_flags({'FLAGS_use_mkldnn': False})
+            paddle.set_flags({'FLAGS_use_mkldnn': False})
 
 
 class TestError(Dy2StTestBase):
-    @test_default_and_pir
+    @test_legacy_and_pt_and_pir
     def test_type_error(self):
         def foo(x):
             out = x + 1
diff --git a/test/dygraph_to_static/test_resnet.py b/test/dygraph_to_static/test_resnet.py
@@ -21,6 +21,8 @@
 import numpy as np
 from dygraph_to_static_utils import (
     Dy2StTestBase,
+    enable_to_static_guard,
+    static_guard,
     test_default_and_pir,
 )
 from predictor_utils import PredictorTools
@@ -247,32 +249,31 @@ def __len__(self):
         return len(self.img)
 
 
-class TestResnet(Dy2StTestBase):
-    def setUp(self):
+class ResNetHelper:
+    def __init__(self):
         self.temp_dir = tempfile.TemporaryDirectory()
 
         self.model_save_dir = os.path.join(self.temp_dir.name, "./inference")
         self.model_save_prefix = os.path.join(
-            self.temp_dir.name, "./inference/resnet_v2"
+            self.temp_dir.name, "./inference/resnet"
         )
         self.model_filename = (
-            "resnet_v2" + paddle.jit.translated_layer.INFER_MODEL_SUFFIX
+            "resnet" + paddle.jit.translated_layer.INFER_MODEL_SUFFIX
         )
         self.params_filename = (
-            "resnet_v2" + paddle.jit.translated_layer.INFER_PARAMS_SUFFIX
+            "resnet" + paddle.jit.translated_layer.INFER_PARAMS_SUFFIX
         )
         self.dy_state_dict_save_path = os.path.join(
-            self.temp_dir.name, "./resnet_v2.dygraph"
+            self.temp_dir.name, "./resnet.dygraph"
         )
 
-    def tearDown(self):
+    def __del__(self):
         self.temp_dir.cleanup()
 
-    def do_train(self, to_static):
+    def train(self, to_static, 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.disable_static(place)
         np.random.seed(SEED)
         paddle.seed(SEED)
         paddle.framework.random._manual_program_seed(SEED)
@@ -284,7 +285,7 @@ def do_train(self, to_static):
             dataset, batch_size=batch_size, drop_last=True
         )
 
-        resnet = paddle.jit.to_static(ResNet())
+        resnet = paddle.jit.to_static(ResNet(), build_strategy=build_strategy)
         optimizer = optimizer_setting(parameter_list=resnet.parameters())
 
         for epoch in range(epoch_num):
@@ -350,59 +351,55 @@ def do_train(self, to_static):
                             self.dy_state_dict_save_path + '.pdparams',
                         )
                     break
-        paddle.enable_static()
 
         return total_loss.numpy()
 
     def predict_dygraph(self, data):
-        paddle.jit.enable_to_static(False)
-        paddle.disable_static(place)
-        resnet = paddle.jit.to_static(ResNet())
+        with enable_to_static_guard(False):
+            resnet = paddle.jit.to_static(ResNet())
 
-        model_dict = paddle.load(self.dy_state_dict_save_path + '.pdparams')
-        resnet.set_dict(model_dict)
-        resnet.eval()
+            model_dict = paddle.load(self.dy_state_dict_save_path + '.pdparams')
+            resnet.set_dict(model_dict)
+            resnet.eval()
 
-        pred_res = resnet(
-            paddle.to_tensor(
-                data=data, dtype=None, place=None, stop_gradient=True
+            pred_res = resnet(
+                paddle.to_tensor(
+                    data=data, dtype=None, place=None, stop_gradient=True
+                )
             )
-        )
 
         ret = pred_res.numpy()
-        paddle.enable_static()
         return ret
 
     def predict_static(self, data):
-        exe = paddle.static.Executor(place)
-        [
-            inference_program,
-            feed_target_names,
-            fetch_targets,
-        ] = paddle.static.load_inference_model(
-            self.model_save_dir,
-            executor=exe,
-            model_filename=self.model_filename,
-            params_filename=self.params_filename,
-        )
+        with static_guard():
+            exe = paddle.static.Executor(place)
+            [
+                inference_program,
+                feed_target_names,
+                fetch_targets,
+            ] = paddle.static.load_inference_model(
+                self.model_save_dir,
+                executor=exe,
+                model_filename=self.model_filename,
+                params_filename=self.params_filename,
+            )
 
-        pred_res = exe.run(
-            inference_program,
-            feed={feed_target_names[0]: data},
-            fetch_list=fetch_targets,
-        )
+            pred_res = exe.run(
+                inference_program,
+                feed={feed_target_names[0]: data},
+                fetch_list=fetch_targets,
+            )
 
-        return pred_res[0]
+            return pred_res[0]
 
     def predict_dygraph_jit(self, data):
-        paddle.disable_static(place)
         resnet = paddle.jit.load(self.model_save_prefix)
         resnet.eval()
 
         pred_res = resnet(data)
 
         ret = pred_res.numpy()
-        paddle.enable_static()
         return ret
 
     def predict_analysis_inference(self, data):
@@ -415,16 +412,21 @@ def predict_analysis_inference(self, data):
         (out,) = output()
         return out
 
+
+class TestResnet(Dy2StTestBase):
+    def setUp(self):
+        self.resnet_helper = ResNetHelper()
+
     def train(self, to_static):
-        paddle.jit.enable_to_static(to_static)
-        return self.do_train(to_static)
+        with enable_to_static_guard(to_static):
+            return self.resnet_helper.train(to_static)
 
     def verify_predict(self):
         image = np.random.random([1, 3, 224, 224]).astype('float32')
-        dy_pre = self.predict_dygraph(image)
-        st_pre = self.predict_static(image)
-        dy_jit_pre = self.predict_dygraph_jit(image)
-        predictor_pre = self.predict_analysis_inference(image)
+        dy_pre = self.resnet_helper.predict_dygraph(image)
+        st_pre = self.resnet_helper.predict_static(image)
+        dy_jit_pre = self.resnet_helper.predict_dygraph_jit(image)
+        predictor_pre = self.resnet_helper.predict_analysis_inference(image)
         np.testing.assert_allclose(
             dy_pre,
             st_pre,
@@ -455,7 +457,7 @@ def test_resnet(self):
             err_msg=f'static_loss: {static_loss} \n dygraph_loss: {dygraph_loss}',
         )
         # TODO(@xiongkun): open after save / load supported in pir.
-        if not paddle.base.framework.use_pir_api():
+        if not paddle.framework.use_pir_api():
             self.verify_predict()
 
     @test_default_and_pir
@@ -474,12 +476,12 @@ def test_resnet_composite(self):
 
     @test_default_and_pir
     def test_in_static_mode_mkldnn(self):
-        paddle.base.set_flags({'FLAGS_use_mkldnn': True})
+        paddle.set_flags({'FLAGS_use_mkldnn': True})
         try:
             if paddle.base.core.is_compiled_with_mkldnn():
                 self.train(to_static=True)
         finally:
-            paddle.base.set_flags({'FLAGS_use_mkldnn': False})
+            paddle.set_flags({'FLAGS_use_mkldnn': False})
 
 
 if __name__ == '__main__':
diff --git a/test/dygraph_to_static/test_seq2seq.py b/test/dygraph_to_static/test_seq2seq.py
@@ -20,9 +20,8 @@
 import numpy as np
 from dygraph_to_static_utils import (
     Dy2StTestBase,
-    ToStaticMode,
-    set_to_static_mode,
     test_legacy_only,
+    test_sot_mgs0_only,
 )
 from seq2seq_dygraph_model import AttentionModel, BaseModel
 from seq2seq_utils import Seq2SeqModelHyperParams, get_data_iter
@@ -239,13 +238,13 @@ def _test_predict(self, attn_model=False):
             msg=f"\npred_dygraph = {pred_dygraph} \npred_static = {pred_static}",
         )
 
-    @set_to_static_mode(ToStaticMode.SOT)
+    @test_sot_mgs0_only
     @test_legacy_only
     def test_base_model(self):
         self._test_train(attn_model=False)
         self._test_predict(attn_model=False)
 
-    @set_to_static_mode(ToStaticMode.SOT)
+    @test_sot_mgs0_only
     @test_legacy_only
     def test_attn_model(self):
         self._test_train(attn_model=True)
diff --git a/test/dygraph_to_static/test_train_step.py b/test/dygraph_to_static/test_train_step.py
@@ -17,7 +17,7 @@
 from functools import partial
 
 import numpy as np
-from dygraph_to_static_utils import Dy2StTestBase
+from dygraph_to_static_utils import Dy2StTestBase, test_ast_only, test_pt_only
 
 import paddle
 
@@ -77,6 +77,8 @@ def get_train_step_losses(self, func, steps):
             losses.append(loss)
         return losses
 
+    @test_ast_only
+    @test_pt_only
     def test_train_step(self):
         reset_seed()
         dygraph_losses = self.get_train_step_losses(
