[PIR / Dy2static] Fix mnist - part 1 (#59447)

2742195759 · changeyoung98 · SigureMo · web-flow · commit df3676e65dd6 · 2023-11-30T11:31:30.000+08:00
---------

Co-authored-by: chenzhiyang &lt;1792266893@qq.com&gt;
Co-authored-by: SigureMo &lt;sigure.qaq@gmail.com&gt;
diff --git a/paddle/phi/infermeta/multiary.cc b/paddle/phi/infermeta/multiary.cc
@@ -675,6 +675,7 @@ void BatchNormInferMeta(const MetaTensor& x,
   }
   if (reserve_space) {
     reserve_space->set_dims({-1});
+    reserve_space->set_dtype(DataType::UINT8);
   }
   y->share_lod(x);
   y->set_dtype(x.dtype());
diff --git a/python/paddle/autograd/ir_backward.py b/python/paddle/autograd/ir_backward.py
@@ -703,8 +703,12 @@ def calc_gradient_helper(outputs, inputs, grad_outputs, no_grad_set):
     )
 
     state.turn_map()
+
     for bwd_op in inverse_sort_op(remove_ops):
-        remove_op(block, bwd_op, state)
+        if bwd_op.result(0) in grad_outputs:
+            continue
+        if bwd_op.result(0).use_empty():
+            remove_op(block, bwd_op, state)
     state.turn_map()
 
     input_grad_map = state.value_to_valuegrad
diff --git a/test/dygraph_to_static/dygraph_to_static_utils.py b/test/dygraph_to_static/dygraph_to_static_utils.py
@@ -404,6 +404,13 @@ def test_sot_with_pir_only(fn):
     return 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.PT | IrMode.PIR)(fn)
+    return fn
+
+
 # NOTE: This is a special decorator for comparing legacy and pt
 def compare_legacy_with_pt(fn):
     @wraps(fn)
diff --git a/test/dygraph_to_static/test_layer_hook.py b/test/dygraph_to_static/test_layer_hook.py
@@ -17,7 +17,10 @@
 import unittest
 
 import numpy as np
-from dygraph_to_static_utils import Dy2StTestBase
+from dygraph_to_static_utils import (
+    Dy2StTestBase,
+    test_legacy_and_pt_and_pir,
+)
 
 import paddle
 
@@ -73,7 +76,8 @@ def train_net(self, to_static=False):
             net = paddle.jit.to_static(net)
         out = net(self.x)
 
-        if to_static:
+        # TODO(xiongkun) save / load unitest.
+        if to_static and not paddle.base.framework.use_pir_api():
             paddle.jit.save(net, self.path)
 
         return float(out)
@@ -83,23 +87,24 @@ def load_train(self):
         out = net(self.x)
         return float(out)
 
+    @test_legacy_and_pt_and_pir
     def test_hook(self):
         dy_out = self.train_net(to_static=False)
         st_out = self.train_net(to_static=True)
-        load_out = self.load_train()
-        print(st_out, dy_out, load_out)
         np.testing.assert_allclose(
             st_out,
             dy_out,
             rtol=1e-05,
             err_msg=f'dygraph_res is {dy_out}\nstatic_res is {st_out}',
         )
-        np.testing.assert_allclose(
-            st_out,
-            load_out,
-            rtol=1e-05,
-            err_msg=f'load_out is {load_out}\nstatic_res is {st_out}',
-        )
+        if not paddle.base.framework.use_pir_api():
+            load_out = self.load_train()
+            np.testing.assert_allclose(
+                st_out,
+                load_out,
+                rtol=1e-05,
+                err_msg=f'load_out is {load_out}\nstatic_res is {st_out}',
+            )
 
 
 if __name__ == "__main__":
diff --git a/test/dygraph_to_static/test_mnist.py b/test/dygraph_to_static/test_mnist.py
@@ -20,15 +20,13 @@
 import numpy as np
 from dygraph_to_static_utils import (
     Dy2StTestBase,
-    test_default_mode_only,
-    test_sot_with_pir_only,
+    test_default_and_pir,
 )
 from predictor_utils import PredictorTools
 
 import paddle
 from paddle import base
 from paddle.base.dygraph import to_variable
-from paddle.base.dygraph.base import switch_to_static_graph
 from paddle.jit.translated_layer import INFER_MODEL_SUFFIX, INFER_PARAMS_SUFFIX
 from paddle.nn import Linear
 from paddle.optimizer import Adam
@@ -163,7 +161,7 @@ def train_static(self):
     def train_dygraph(self):
         return self.train(to_static=False)
 
-    @test_default_mode_only
+    @test_default_and_pir
     def test_mnist_to_static(self):
         dygraph_loss = self.train_dygraph()
         static_loss = self.train_static()
@@ -174,7 +172,7 @@ def test_mnist_to_static(self):
             err_msg=f'dygraph is {dygraph_loss}\n static_res is \n{static_loss}',
         )
 
-    @test_sot_with_pir_only
+    @test_default_and_pir
     def test_mnist_declarative_cpu_vs_mkldnn(self):
         dygraph_loss_cpu = self.train_dygraph()
         paddle.set_flags({'FLAGS_use_mkldnn': True})
@@ -239,14 +237,16 @@ def train(self, to_static=False):
                     prediction, acc, avg_loss = mnist(img, label)
                     loss_data.append(float(avg_loss))
                     # new save load check
-                    self.check_jit_save_load(
-                        mnist,
-                        [dy_x_data],
-                        [img, label],
-                        to_static,
-                        prediction,
-                        [img.name],
-                    )
+                    # TODO(@xiongkun): enable this after new save load is supported in pir.
+                    if not paddle.framework.use_pir_api():
+                        self.check_jit_save_load(
+                            mnist,
+                            [dy_x_data],
+                            [img, label],
+                            to_static,
+                            prediction,
+                            [img.name],
+                        )
                     break
         return loss_data
 
@@ -298,7 +298,6 @@ def check_jit_save_load(
                 gt_out.numpy(), predictor_infer_out, rtol=1e-05
             )
 
-    @switch_to_static_graph
     def jit_load_and_run_inference_static(
         self, model_path, model_filename, params_filename, inputs
     ):
@@ -320,6 +319,7 @@ def jit_load_and_run_inference_static(
             feed=dict(zip(feed_target_names, inputs)),
             fetch_list=fetch_targets,
         )
+        paddle.disable_static()
 
         return np.array(results[0])
 
diff --git a/test/dygraph_to_static/test_word2vec.py b/test/dygraph_to_static/test_word2vec.py
@@ -17,7 +17,10 @@
 import unittest
 
 import numpy as np
-from dygraph_to_static_utils import Dy2StTestBase
+from dygraph_to_static_utils import (
+    Dy2StTestBase,
+    test_legacy_and_pt_and_pir,
+)
 
 import paddle
 from paddle import base
@@ -317,6 +320,7 @@ 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)

Original file line number	Diff line number	Diff line change
`@@ -675,6 +675,7 @@ void BatchNormInferMeta(const MetaTensor& x,`
`675`	`675`	`}`
`676`	`676`	`if (reserve_space) {`
`677`	`677`	`reserve_space->set_dims({-1});`
	`678`	`+ reserve_space->set_dtype(DataType::UINT8);`
`678`	`679`	`}`
`679`	`680`	`y->share_lod(x);`
`680`	`681`	`y->set_dtype(x.dtype());`