[Dy2St][NO.11] pir dy2st unittest fix test_save_inference_model - Part 3 (#60211)

gouzil · web-flow · commit d6f6402da299 · 2023-12-21T18:31:15.000+08:00
diff --git a/python/paddle/nn/functional/common.py b/python/paddle/nn/functional/common.py
@@ -1960,9 +1960,9 @@ def linear(x, weight, bias=None, name=None):
         return _C_ops.linear(x, weight, bias)
 
     elif in_pir_mode():
-        out = paddle._pir_ops.matmul(x, weight, False, False)
+        out = _C_ops.matmul(x, weight, False, False)
         if bias is not None:
-            return paddle._pir_ops.add(out, bias)
+            return _C_ops.add(out, bias)
         else:
             return out
     else:
diff --git a/test/dygraph_to_static/test_save_inference_model.py b/test/dygraph_to_static/test_save_inference_model.py
@@ -21,28 +21,33 @@
     Dy2StTestBase,
     compare_legacy_with_pt,
     test_ast_only,
+    test_legacy_and_pt_and_pir,
 )
 
 import paddle
 from paddle import base
 from paddle.autograd import PyLayer
-from paddle.jit.api import to_static
+from paddle.framework import use_pir_api
 from paddle.jit.dy2static.partial_program import partial_program_from
+from paddle.jit.dy2static.pir_partial_program import (
+    partial_program_from as pir_partial_program_from,
+)
 from paddle.jit.translated_layer import INFER_MODEL_SUFFIX, INFER_PARAMS_SUFFIX
 
 SEED = 2020
 
 np.random.seed(SEED)
 
-place = base.CUDAPlace(0) if base.is_compiled_with_cuda() else base.CPUPlace()
+place = (
+    paddle.CUDAPlace(0) if paddle.is_compiled_with_cuda() else paddle.CPUPlace()
+)
 
 
 class SimpleFcLayer(paddle.nn.Layer):
     def __init__(self, fc_size):
         super().__init__()
         self._linear = paddle.nn.Linear(fc_size, fc_size)
 
-    @to_static(full_graph=True)
     def forward(self, x):
         y = self._linear(x)
         z = self._linear(y)
@@ -69,7 +74,6 @@ def __init__(self, fc_size):
         super().__init__()
         self._linear = paddle.nn.Linear(fc_size, fc_size)
 
-    @to_static(full_graph=True)
     def forward(self, x):
         y = self._linear(x)
         out = cus_tanh.apply(y)
@@ -88,28 +92,29 @@ def tearDown(self):
     def test_save_inference_model(self):
         fc_size = 20
         x_data = np.random.random((fc_size, fc_size)).astype('float32')
-        with base.dygraph.guard(place):
-            base.default_startup_program().random_seed = SEED
-            base.default_main_program().random_seed = SEED
-
-            x = base.dygraph.to_variable(x_data)
-            layer = SimpleFcLayer(fc_size)
-            adam = paddle.optimizer.SGD(
-                learning_rate=0.1, parameters=layer.parameters()
-            )
+        base.default_startup_program().random_seed = SEED
+        base.default_main_program().random_seed = SEED
 
-            for i in range(5):
-                loss, pred = layer(x)
-                loss.backward()
-                adam.minimize(loss)
-                layer.clear_gradients()
-            # test for saving model in dygraph.guard
-            infer_model_prefix = os.path.join(
-                self.temp_dir.name, "test_dy2stat_inference_in_guard/model"
-            )
-            infer_model_dir = os.path.join(
-                self.temp_dir.name, "test_dy2stat_inference_in_guard"
-            )
+        x = paddle.to_tensor(x_data)
+        layer = paddle.jit.to_static(SimpleFcLayer(fc_size))
+        adam = paddle.optimizer.SGD(
+            learning_rate=0.1, parameters=layer.parameters()
+        )
+
+        for i in range(5):
+            loss, pred = layer(x)
+            loss.backward()
+            adam.minimize(loss)
+            layer.clear_gradients()
+        # test for saving model in dygraph.guard
+        infer_model_prefix = os.path.join(
+            self.temp_dir.name, "test_dy2stat_inference_in_guard/model"
+        )
+        infer_model_dir = os.path.join(
+            self.temp_dir.name, "test_dy2stat_inference_in_guard"
+        )
+        # TODO(pir-save-load): Fix this after we support save/load in PIR
+        if not use_pir_api():
             paddle.jit.save(
                 layer=layer,
                 path=infer_model_prefix,
@@ -118,25 +123,26 @@ def test_save_inference_model(self):
             )
             # Check the correctness of the inference
             dygraph_out, _ = layer(x)
-        self.check_save_inference_model(layer, [x_data], dygraph_out.numpy())
-        self.check_save_inference_model(
-            layer, [x_data], dygraph_out.numpy(), fetch=[loss]
-        )
-        self.check_save_inference_model(
-            layer, [x_data], dygraph_out.numpy(), feed=[x]
-        )
+            self.check_save_inference_model(
+                layer, [x_data], dygraph_out.numpy()
+            )
+            self.check_save_inference_model(
+                layer, [x_data], dygraph_out.numpy(), fetch=[loss]
+            )
+            self.check_save_inference_model(
+                layer, [x_data], dygraph_out.numpy(), feed=[x]
+            )
 
     @test_ast_only
     def test_save_pylayer_model(self):
         fc_size = 20
         x_data = np.random.random((fc_size, fc_size)).astype('float32')
-        paddle.base.framework._set_expected_place(place)
+        paddle.framework._set_expected_place(place)
 
         base.default_startup_program().random_seed = SEED
         base.default_main_program().random_seed = SEED
-        paddle.disable_static()
-        x = base.dygraph.to_variable(x_data)
-        layer = SimplePyLayerNet(fc_size)
+        x = paddle.to_tensor(x_data)
+        layer = paddle.jit.to_static(SimplePyLayerNet(fc_size))
         adam = paddle.optimizer.SGD(
             learning_rate=0.1, parameters=layer.parameters()
         )
@@ -150,25 +156,27 @@ def test_save_pylayer_model(self):
         infer_model_prefix = os.path.join(
             self.temp_dir.name, "test_dy2stat_inference_in_guard/model_pylayer"
         )
-        paddle.jit.save(
-            layer=layer,
-            path=infer_model_prefix,
-            input_spec=[x],
-            output_spec=[pred],
-        )
-        # Check the correctness of the inference
-        loss_out, _ = layer(x)
+        # TODO(pir-save-load): Fix this after we support save/load in PIR
+        if not use_pir_api():
+            paddle.jit.save(
+                layer=layer,
+                path=infer_model_prefix,
+                input_spec=[x],
+                output_spec=[pred],
+            )
+            # Check the correctness of the inference
+            loss_out, _ = layer(x)
 
-        loss_out_numpy = float(loss_out)
-        self.check_save_inference_model(
-            layer, [x_data], loss_out_numpy, enable_pir=False
-        )
-        self.check_save_inference_model(
-            layer, [x_data], loss_out_numpy, fetch=[loss], enable_pir=False
-        )
-        self.check_save_inference_model(
-            layer, [x_data], loss_out_numpy, feed=[x], enable_pir=False
-        )
+            loss_out_numpy = float(loss_out)
+            self.check_save_inference_model(
+                layer, [x_data], loss_out_numpy, enable_pir=False
+            )
+            self.check_save_inference_model(
+                layer, [x_data], loss_out_numpy, fetch=[loss], enable_pir=False
+            )
+            self.check_save_inference_model(
+                layer, [x_data], loss_out_numpy, feed=[x], enable_pir=False
+            )
 
     def check_save_inference_model(
         self, model, inputs, gt_out, feed=None, fetch=None, enable_pir=True
@@ -224,30 +232,36 @@ def load_and_run_inference(
             fetch_list=fetch_targets,
         )
 
+        paddle.disable_static()
         return np.array(results[0])
 
 
 class TestPartialProgramRaiseError(Dy2StTestBase):
     @test_ast_only
+    @test_legacy_and_pt_and_pir
     def test_param_type(self):
         x_data = np.random.random((20, 20)).astype('float32')
 
-        with base.dygraph.guard(base.CPUPlace()):
-            net = SimpleFcLayer(20)
-            x = base.dygraph.to_variable(x_data)
-            out = net(x)
+        net = paddle.jit.to_static(SimpleFcLayer(20))
+        x = paddle.to_tensor(x_data)
+        out = net(x)
 
-            program_cache = net.forward.program_cache
-            _, (concrete_program, _) = program_cache.last()
+        program_cache = net.forward.program_cache
+        _, (concrete_program, _) = program_cache.last()
 
-            params = concrete_program.parameters
+        params = concrete_program.parameters
 
-            concrete_program.parameters = params[0]
-            # TypeError: Type of self._params should be list or tuple,
-            # but received <class 'paddle.base.framework.EagerParamBase'>.
-            with self.assertRaises(TypeError):
+        concrete_program.parameters = params[0]
+        # TypeError: Type of self._params should be list or tuple,
+        # but received <class 'paddle.base.framework.EagerParamBase'>.
+        with self.assertRaises(TypeError):
+            if use_pir_api():
+                pir_partial_program_from(concrete_program)
+            else:
                 partial_program_from(concrete_program)
 
+        # Under PIR, params are tuples and cannot be modified
+        if not use_pir_api():
             params[0] = "linear.w.0"
             concrete_program.parameters = params
             # TypeError: Type of self._params[0] should be framework.EagerParamBase,
