[NewIR]Enable dy2st new ir test part2 (#56061)

* support new ir dy2st

* revert code

* skip test with stride

* chang import file

* enable dy2st new ir test part2

* update

* suppor train step

* fix bug

* fix bug

* update

test/dygraph_to_static/test_bmn.py

@@ -18,7 +18,7 @@ import tempfile
 import unittest
 
 import numpy as np
-from dygraph_to_static_util import dy2static_unittest
+from dygraph_to_static_util import dy2static_unittest, test_with_new_ir
 from predictor_utils import PredictorTools
 
 import paddle
@@ -751,6 +751,21 @@ class TestTrain(unittest.TestCase):
                         break
             return np.array(loss_data)
 
+    @test_with_new_ir
+    def test_train_new_ir(self):
+        static_res = self.train_bmn(self.args, self.place, to_static=True)
+        dygraph_res = self.train_bmn(self.args, self.place, to_static=False)
+        np.testing.assert_allclose(
+            dygraph_res,
+            static_res,
+            rtol=1e-05,
+            err_msg='dygraph_res: {},\n static_res: {}'.format(
+                dygraph_res[~np.isclose(dygraph_res, static_res)],
+                static_res[~np.isclose(dygraph_res, static_res)],
+            ),
+            atol=1e-8,
+        )
+
     def test_train(self):
         static_res = self.train_bmn(self.args, self.place, to_static=True)
         dygraph_res = self.train_bmn(self.args, self.place, to_static=False)

test/dygraph_to_static/test_cycle_gan.py

@@ -36,6 +36,8 @@ from paddle import fluid
 # Use GPU:0 to elimate the influence of other tasks.
 os.environ["CUDA_VISIBLE_DEVICES"] = "1"
 
+from dygraph_to_static_util import test_and_compare_with_new_ir
+
 import paddle
 from paddle.fluid.dygraph import to_variable
 from paddle.jit.api import to_static
@@ -694,6 +696,7 @@ class TestCycleGANModel(unittest.TestCase):
         out = train(self.args, to_static)
         return out
 
+    @test_and_compare_with_new_ir(False)
     def test_train(self):
         st_out = self.train(to_static=True)
         dy_out = self.train(to_static=False)

test/dygraph_to_static/test_mnist_amp.py

@@ -16,6 +16,7 @@ import unittest
 from time import time
 
 import numpy as np
+from dygraph_to_static_util import test_and_compare_with_new_ir
 from test_mnist import MNIST, SEED, TestMNIST
 
 import paddle
@@ -32,6 +33,7 @@ class TestAMP(TestMNIST):
     def train_dygraph(self):
         return self.train(to_static=False)
 
+    @test_and_compare_with_new_ir(False)
     def test_mnist_to_static(self):
         dygraph_loss = self.train_dygraph()
         static_loss = self.train_static()

test/dygraph_to_static/test_mnist_pure_fp16.py

@@ -16,6 +16,7 @@ import unittest
 from time import time
 
 import numpy as np
+from dygraph_to_static_util import test_and_compare_with_new_ir
 from test_mnist import MNIST, SEED, TestMNIST
 
 import paddle
@@ -31,6 +32,7 @@ class TestPureFP16(TestMNIST):
     def train_dygraph(self):
         return self.train(to_static=False)
 
+    @test_and_compare_with_new_ir(False)
     def test_mnist_to_static(self):
         if paddle.fluid.is_compiled_with_cuda():
             dygraph_loss = self.train_dygraph()

test/dygraph_to_static/test_resnet_amp.py

@@ -16,6 +16,7 @@ import time
 import unittest
 
 import numpy as np
+from dygraph_to_static_util import test_and_compare_with_new_ir
 from test_resnet import SEED, ResNet, optimizer_setting
 
 import paddle
@@ -117,6 +118,7 @@ class TestResnet(unittest.TestCase):
         paddle.jit.enable_to_static(to_static)
         return train(to_static)
 
+    @test_and_compare_with_new_ir(False)
     def test_resnet(self):
         static_loss = self.train(to_static=True)
         dygraph_loss = self.train(to_static=False)

test/dygraph_to_static/test_resnet_pure_fp16.py

@@ -16,6 +16,7 @@ import time
 import unittest
 
 import numpy as np
+from dygraph_to_static_util import test_and_compare_with_new_ir
 from test_resnet import SEED, ResNet, optimizer_setting
 
 import paddle
@@ -120,6 +121,7 @@ class TestResnet(unittest.TestCase):
         build_strategy.enable_inplace = False
         return train(to_static, build_strategy)
 
+    @test_and_compare_with_new_ir(False)
     def test_resnet(self):
         if fluid.is_compiled_with_cuda():
             static_loss = self.train(to_static=True)

test/dygraph_to_static/test_train_step.py

@@ -17,7 +17,10 @@ import unittest
 from functools import partial
 
 import numpy as np
-from dygraph_to_static_util import enable_fallback_guard
+from dygraph_to_static_util import (
+    enable_fallback_guard,
+    test_and_compare_with_new_ir,
+)
 
 import paddle
 
@@ -77,6 +80,7 @@ class TestTrainStepTinyModel(unittest.TestCase):
             losses.append(loss)
         return losses
 
+    @test_and_compare_with_new_ir(False)
     def test_train_step(self):
         reset_seed()
         dygraph_losses = self.get_train_step_losses(

test/dygraph_to_static/test_tsm.py

@@ -19,6 +19,7 @@ import sys
 import unittest
 
 import numpy as np
+from dygraph_to_static_util import test_and_compare_with_new_ir
 from tsm_config_utils import merge_configs, parse_config, print_configs
 
 import paddle
@@ -382,6 +383,7 @@ def train(args, fake_data_reader, to_static):
 
 
 class TestTsm(unittest.TestCase):
+    @test_and_compare_with_new_ir(False)
     def test_dygraph_static_same_loss(self):
         if fluid.is_compiled_with_cuda():
             fluid.set_flags({"FLAGS_cudnn_deterministic": True})

test/dygraph_to_static/test_word2vec.py

@@ -17,6 +17,7 @@ import random
 import unittest
 
 import numpy as np
+from dygraph_to_static_util import test_and_compare_with_new_ir
 
 import paddle
 from paddle import fluid
@@ -320,6 +321,7 @@ def train(to_static):
 
 
 class TestWord2Vec(unittest.TestCase):
+    @test_and_compare_with_new_ir(False)
     def test_dygraph_static_same_loss(self):
         dygraph_loss = train(to_static=False)
         static_loss = train(to_static=True)

test/dygraph_to_static/test_yolov3.py

@@ -17,6 +17,7 @@ import time
 import unittest
 
 import numpy as np
+from dygraph_to_static_util import test_and_compare_with_new_ir
 from yolov3 import YOLOv3, cfg
 
 import paddle
@@ -166,6 +167,7 @@ def train(to_static):
 
 
 class TestYolov3(unittest.TestCase):
+    @test_and_compare_with_new_ir(False)
     def test_dygraph_static_same_loss(self):
         dygraph_loss = train(to_static=False)
         static_loss = train(to_static=True)