[Dy2St][NO.13] pir dy2st unittest fix test_seq2seq - Part 4 (#60454)

gouzil · SigureMo · web-flow · commit 0d1009c9bd62 · 2024-02-21T20:42:34.000+08:00
---------

Co-authored-by: SigureMo &lt;sigure.qaq@gmail.com&gt;
diff --git a/test/dygraph_to_static/CMakeLists.txt b/test/dygraph_to_static/CMakeLists.txt
@@ -31,6 +31,7 @@ if(NOT WITH_GPU)
   list(REMOVE_ITEM TEST_OPS test_bert)
   list(REMOVE_ITEM TEST_OPS test_transformer)
   list(REMOVE_ITEM TEST_OPS test_mobile_net)
+  list(REMOVE_ITEM TEST_OPS test_seq2seq)
 endif()
 
 foreach(TEST_OP ${TEST_OPS})
@@ -40,7 +41,6 @@ endforeach()
 set_tests_properties(test_se_resnet PROPERTIES TIMEOUT 900)
 set_tests_properties(test_yolov3 PROPERTIES TIMEOUT 900 LABELS
                                             "RUN_TYPE=EXCLUSIVE")
-set_tests_properties(test_seq2seq PROPERTIES TIMEOUT 420)
 set_tests_properties(test_cycle_gan PROPERTIES TIMEOUT 150)
 set_tests_properties(test_basic_api_transformation PROPERTIES TIMEOUT 240)
 set_tests_properties(test_reinforcement_learning PROPERTIES TIMEOUT 120)
@@ -62,6 +62,7 @@ if(WITH_GPU)
   set_tests_properties(test_transformer PROPERTIES TIMEOUT 240)
   set_tests_properties(test_mobile_net PROPERTIES TIMEOUT 240)
   set_tests_properties(test_resnet PROPERTIES TIMEOUT 240)
+  set_tests_properties(test_seq2seq PROPERTIES TIMEOUT 900)
 endif()
 
 # Legacy IR only tests for dygraph_to_static
diff --git a/test/dygraph_to_static/test_seq2seq.py b/test/dygraph_to_static/test_seq2seq.py
@@ -21,19 +21,15 @@
 from dygraph_to_static_utils import (
     Dy2StTestBase,
     enable_to_static_guard,
-    test_legacy_only,
-    test_sot_mgs0_only,
+    test_default_and_pir,
 )
 from seq2seq_dygraph_model import AttentionModel, BaseModel
 from seq2seq_utils import Seq2SeqModelHyperParams, get_data_iter
 
 import paddle
-from paddle import base
+from paddle.base.framework import unique_name
 from paddle.nn import ClipGradByGlobalNorm
 
-place = (
-    paddle.CUDAPlace(0) if paddle.is_compiled_with_cuda() else paddle.CPUPlace()
-)
 STEP_NUM = 10
 PRINT_STEP = 2
 
@@ -51,7 +47,7 @@ def prepare_input(batch):
 
 
 def train(args, attn_model=False):
-    with base.dygraph.guard(place):
+    with unique_name.guard():
         paddle.seed(2020)
 
         if attn_model:
@@ -141,52 +137,49 @@ def train(args, attn_model=False):
 
 
 def infer(args, attn_model=False):
-    with base.dygraph.guard(place):
-        if attn_model:
-            model = paddle.jit.to_static(
-                AttentionModel(
-                    args.hidden_size,
-                    args.src_vocab_size,
-                    args.tar_vocab_size,
-                    args.batch_size,
-                    beam_size=args.beam_size,
-                    num_layers=args.num_layers,
-                    init_scale=args.init_scale,
-                    dropout=0.0,
-                    mode='beam_search',
-                )
+    if attn_model:
+        model = paddle.jit.to_static(
+            AttentionModel(
+                args.hidden_size,
+                args.src_vocab_size,
+                args.tar_vocab_size,
+                args.batch_size,
+                beam_size=args.beam_size,
+                num_layers=args.num_layers,
+                init_scale=args.init_scale,
+                dropout=0.0,
+                mode='beam_search',
             )
-        else:
-            model = paddle.jit.to_static(
-                BaseModel(
-                    args.hidden_size,
-                    args.src_vocab_size,
-                    args.tar_vocab_size,
-                    args.batch_size,
-                    beam_size=args.beam_size,
-                    num_layers=args.num_layers,
-                    init_scale=args.init_scale,
-                    dropout=0.0,
-                    mode='beam_search',
-                )
+        )
+    else:
+        model = paddle.jit.to_static(
+            BaseModel(
+                args.hidden_size,
+                args.src_vocab_size,
+                args.tar_vocab_size,
+                args.batch_size,
+                beam_size=args.beam_size,
+                num_layers=args.num_layers,
+                init_scale=args.init_scale,
+                dropout=0.0,
+                mode='beam_search',
             )
-
-        model_path = (
-            args.attn_model_path if attn_model else args.base_model_path
         )
-        state_dict = paddle.load(model_path + '.pdparams')
-        model.set_dict(state_dict)
-        model.eval()
-        train_data_iter = get_data_iter(args.batch_size, mode='infer')
-        for batch_id, batch in enumerate(train_data_iter):
-            input_data_feed, word_num = prepare_input(batch)
-            input_data_feed = [
-                paddle.to_tensor(np_inp) for np_inp in input_data_feed
-            ]
-            outputs = paddle.jit.to_static(model.beam_search)(input_data_feed)
-            break
 
-        return outputs.numpy()
+    model_path = args.attn_model_path if attn_model else args.base_model_path
+    state_dict = paddle.load(model_path + '.pdparams')
+    model.set_dict(state_dict)
+    model.eval()
+    train_data_iter = get_data_iter(args.batch_size, mode='infer')
+    for batch_id, batch in enumerate(train_data_iter):
+        input_data_feed, word_num = prepare_input(batch)
+        input_data_feed = [
+            paddle.to_tensor(np_inp) for np_inp in input_data_feed
+        ]
+        outputs = paddle.jit.to_static(model.beam_search)(input_data_feed)
+        break
+
+    return outputs.numpy()
 
 
 class TestSeq2seq(Dy2StTestBase):
@@ -237,14 +230,12 @@ def _test_predict(self, attn_model=False):
             msg=f"\npred_dygraph = {pred_dygraph} \npred_static = {pred_static}",
         )
 
-    @test_sot_mgs0_only
-    @test_legacy_only
+    @test_default_and_pir
     def test_base_model(self):
         self._test_train(attn_model=False)
         self._test_predict(attn_model=False)
 
-    @test_sot_mgs0_only
-    @test_legacy_only
+    @test_default_and_pir
     def test_attn_model(self):
         self._test_train(attn_model=True)
         # TODO(liym27): add predict
