【Fix PIR JIT SaveLoad Unittest No.12,22-23】modify test_se_resnet.py , test_build_strategy.py (#64362)

xiaoguoguo626807 · web-flow · commit 5e8ab2f19647 · 2024-05-17T10:55:29.000+08:00
* modify test_modbile_net.py

* delete print

* modify ci

* modify se_resnet

* add change

* add test_input_sepc
diff --git a/paddle/fluid/pybind/eager_utils.cc b/paddle/fluid/pybind/eager_utils.cc
@@ -2018,6 +2018,9 @@ paddle::Tensor CreateTensorFromValue(const pir::Value& value) {
 
   auto dims = phi::vectorize(GetValueDims(value));
   auto ddims = phi::make_ddim(dims);
+  if (HasValueName(value)) {
+    tensor.set_name(GetValueName(value));
+  }
   auto autograd_meta = egr::EagerUtils::autograd_meta(&tensor);
   autograd_meta->SetPersistable(false);
   autograd_meta->SetStopGradient(GetValueBoolAttr(value, kAttrStopGradients));
diff --git a/paddle/fluid/pybind/pir.h b/paddle/fluid/pybind/pir.h
@@ -25,5 +25,7 @@ using pir::Value;
 void BindPir(pybind11::module *m);
 const phi::DDim &GetValueDims(Value value);
 bool GetValueBoolAttr(Value value, const std::string &attr_name);
+std::string GetValueName(Value value);
+bool HasValueName(const Value &value);
 }  // namespace pybind
 }  // namespace paddle
diff --git a/python/paddle/jit/api.py b/python/paddle/jit/api.py
@@ -534,7 +534,9 @@ def _get_output_vars(outputs, output_spec, with_hook=False):
     output_spec_is_not_value_error = (
         "tensor `%s` is not support in pir mode, "
         "because pir value has no name sometimes, especially as ouptut,"
-        " so we can't check tensor's name with output var name"
+        "so we can't check tensor's name with output var name, please"
+        "change as pir.value(to_static layer's output)"
+        "or int(the position of to_static layer's output)"
     )
     if output_spec and with_hook:
         raise RuntimeError(
@@ -549,31 +551,43 @@ def _get_output_vars(outputs, output_spec, with_hook=False):
                 result_list.append(var)
 
         if output_spec is not None:
-            if len(output_spec) == len(result_list):
+            output_size = len(result_list)
+            if len(output_spec) == output_size:
                 for var in output_spec:
-                    if not isinstance(var, paddle.pir.Value):
+                    if not isinstance(var, paddle.pir.Value, int):
                         warnings.warn(output_spec_is_not_value_error % var.name)
                     else:
                         if var not in ValueSet(result_list):
                             warnings.warn(name_no_exists_error % var.name)
             else:
                 result_set = ValueSet(result_list)
-                result_list = []
+                part_result_list = []
                 for var in output_spec:
-                    if not isinstance(var, paddle.pir.Value):
-                        raise ValueError(
-                            output_spec_is_not_value_error % var.name
-                        )
-                    else:
+                    if isinstance(var, paddle.pir.Value):
                         if var not in result_set:
                             raise ValueError(name_no_exists_error % var.name)
                         else:
-                            result_list.append(var)
+                            part_result_list.append(var)
+                    elif isinstance(var, int):
+                        if var >= output_size:
+                            raise ValueError(
+                                "position %d should smaller than output's size % d",
+                                var,
+                                output_size,
+                            )
+                        else:
+                            part_result_list.append(result_list[var])
+
+                    else:
+                        raise ValueError(
+                            output_spec_is_not_value_error % var.name
+                        )
 
+                return part_result_list
     else:
         output_vars_dict = OrderedDict()
         for var in paddle.utils.flatten(outputs):
-            if isinstance(var, Variable):
+            if isinstance(var, (Variable)):
                 output_vars_dict[var.name] = var
         if output_spec is None:
             result_list = list(output_vars_dict.values())
@@ -821,10 +835,14 @@ def save(layer, path, input_spec=None, **configs):
             recommend using these configurations, they may be removed in the future. If not necessary,
             DO NOT use them. Default None.
             The following options are currently supported:
-            (1) output_spec (list[Tensor]): Selects the output targets of the saved model.
+            (1) output_spec (list[Tensor|Value|int]): Selects the output targets of the saved model,
             By default, all return variables of original Layer's forward method are kept as the
             output of the saved model. If the provided ``output_spec`` list is not all output variables,
             the saved model will be pruned according to the given ``output_spec`` list.
+            in pir mode, Tensor is not supported, because value has no name in most cases,
+            which can't be used to judge which tensor corresponds to which value; the value can't be found
+            if the saved program is not the same as the program that includes output_spec, so we need to
+            use the position of the output.
 
     Returns:
         None
@@ -1197,6 +1215,10 @@ def save(layer, path, input_spec=None, **configs):
         # the rule is like [ Get input variables name ]. For output var,
         # we only support Tensor spec, and actually, we only need the
         # var name of output, and we don't recommended to use output_spec
+        # NOTE(Ruting): in pir mode, Tensor is not supported, because value has no name in most cases,
+        # which can't be used to judge which tensor corresponds to which value; the value can't be found
+        # if the saved program is not the same as the program that includes output_spec, so we need to
+        # use the position of the output.
 
         output_vars = _get_output_vars(
             concrete_program.outputs, configs.output_spec, with_hook
diff --git a/python/paddle/static/pir_io.py b/python/paddle/static/pir_io.py
@@ -183,7 +183,7 @@ def _inference_optimize(program, prune_read_op=True):
 
     # remove all readers and the read_op if exist
     if prune_read_op:
-        warnings.warn("WAINING : need to remove the read_op, please add code ")
+        pass
 
     # change all `is_test` attributes to True
     for block in program.blocks:
diff --git a/test/deprecated/legacy_test/test_input_spec.py b/test/deprecated/legacy_test/test_input_spec.py
@@ -22,6 +22,7 @@
 from paddle.base import core
 from paddle.base.framework import convert_np_dtype_to_dtype_
 from paddle.jit.dy2static.utils import _compatible_non_tensor_spec
+from paddle.pir_utils import test_with_dygraph_pir
 from paddle.static import InputSpec
 
 
@@ -167,18 +168,22 @@ def tearDown(self):
     def setUpClass(cls):
         paddle.disable_static()
 
+    @test_with_dygraph_pir
     def test_non_tensor_bool(self):
         specs = [self.x_spec, False]
         self.check_result(specs, 'bool')
 
+    @test_with_dygraph_pir
     def test_non_tensor_str(self):
         specs = [self.x_spec, True, "xxx"]
         self.check_result(specs, 'str')
 
+    @test_with_dygraph_pir
     def test_non_tensor_int(self):
         specs = [self.x_spec, True, "bn", 10]
         self.check_result(specs, 'int')
 
+    @test_with_dygraph_pir
     def test_non_tensor_list(self):
         specs = [self.x_spec, False, "bn", -10, [4]]
         self.check_result(specs, 'list')
@@ -213,6 +218,7 @@ def check_result(self, specs, path):
 
         np.testing.assert_allclose(st_out, load_out, rtol=1e-05)
 
+    @test_with_dygraph_pir
     def test_spec_compatible(self):
         net = NetWithNonTensorSpec(self.in_num, self.out_num)
 
@@ -270,6 +276,7 @@ def setUp(self):
     def setUpClass(cls):
         paddle.disable_static()
 
+    @test_with_dygraph_pir
     def test_non_tensor_with_prune(self):
         specs = [self.x_spec, self.y_spec, True]
         path = os.path.join(self.temp_dir.name, './net_non_tensor_prune_')
@@ -295,11 +302,16 @@ def test_non_tensor_with_prune(self):
 
         # jit.save and jit.load with prune y and loss
         prune_specs = [self.x_spec, True]
+        if paddle.framework.use_pir_api():
+            output_spec = [0]
+        else:
+            output_spec = [st_out]
+
         paddle.jit.save(
             net,
             path,
             prune_specs,
-            output_spec=[st_out],
+            output_spec=output_spec,
             input_names_after_prune=[self.x_spec.name],
         )
         load_net = paddle.jit.load(path)
diff --git a/test/dygraph_to_static/test_build_strategy.py b/test/dygraph_to_static/test_build_strategy.py
@@ -24,6 +24,7 @@
 from test_resnet import ResNetHelper
 
 import paddle
+from paddle.framework import use_pir_api
 
 
 class TestResnetWithPass(Dy2StTestBase):
@@ -46,7 +47,6 @@ def verify_predict(self):
         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,
@@ -59,12 +59,14 @@ def verify_predict(self):
             rtol=1e-05,
             err_msg=f'dy_jit_pre:\n {dy_jit_pre}\n, st_pre: \n{st_pre}.',
         )
-        np.testing.assert_allclose(
-            predictor_pre,
-            st_pre,
-            rtol=1e-05,
-            err_msg=f'predictor_pre:\n {predictor_pre}\n, st_pre: \n{st_pre}.',
-        )
+        if not use_pir_api():
+            predictor_pre = self.resnet_helper.predict_analysis_inference(image)
+            np.testing.assert_allclose(
+                predictor_pre,
+                st_pre,
+                rtol=1e-05,
+                err_msg=f'predictor_pre:\n {predictor_pre}\n, st_pre: \n{st_pre}.',
+            )
 
     @test_default_and_pir
     def test_resnet(self):
@@ -76,9 +78,7 @@ def test_resnet(self):
             rtol=1e-05,
             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():
-            self.verify_predict()
+        self.verify_predict()
 
     @test_default_and_pir
     def test_in_static_mode_mkldnn(self):
diff --git a/test/dygraph_to_static/test_se_resnet.py b/test/dygraph_to_static/test_se_resnet.py
@@ -29,7 +29,8 @@
 
 import paddle
 from paddle import base
-from paddle.jit.api import to_static
+from paddle.framework import use_pir_api
+from paddle.jit.pir_translated_layer import PIR_INFER_MODEL_SUFFIX
 from paddle.jit.translated_layer import INFER_MODEL_SUFFIX, INFER_PARAMS_SUFFIX
 from paddle.nn import BatchNorm, Linear
 from paddle.static import InputSpec
@@ -322,7 +323,6 @@ def __init__(self, layers=50, class_dim=102):
             ),
         )
 
-    @to_static(full_graph=True)
     def forward(self, inputs, label):
         if self.layers == 50 or self.layers == 101:
             y = self.conv0(inputs)
@@ -366,6 +366,7 @@ def setUp(self):
             self.temp_dir.name, "inference/se_resnet"
         )
         self.model_filename = "se_resnet" + INFER_MODEL_SUFFIX
+        self.pir_model_filename = "se_resnet" + PIR_INFER_MODEL_SUFFIX
         self.params_filename = "se_resnet" + INFER_PARAMS_SUFFIX
         self.dy_state_dict_save_path = os.path.join(
             self.temp_dir.name, "se_resnet.dygraph"
@@ -381,6 +382,7 @@ def train(self, train_reader, to_static):
             paddle.seed(SEED)
             paddle.framework.random._manual_program_seed(SEED)
             se_resnext = SeResNeXt()
+            se_resnext = paddle.jit.to_static(se_resnext, full_graph=True)
             optimizer = optimizer_setting(
                 train_parameters, se_resnext.parameters()
             )
@@ -450,15 +452,16 @@ def train(self, train_reader, to_static):
 
                     step_idx += 1
                     if step_idx == STEP_NUM:
-                        # TODO(@xiongkun): open after save / load supported in pir.
-                        if (
-                            to_static
-                            and not paddle.base.framework.use_pir_api()
-                        ):
+                        if to_static:
+                            if use_pir_api():
+                                output_spec = [0]
+                            else:
+                                output_spec = [pred]
+
                             paddle.jit.save(
                                 se_resnext,
                                 self.model_save_prefix,
-                                output_spec=[pred],
+                                output_spec=output_spec,
                                 input_names_after_prune=['x'],
                                 input_spec=[
                                     InputSpec(
@@ -501,6 +504,11 @@ def predict_dygraph(self, data):
 
     def predict_static(self, data):
         paddle.enable_static()
+        if use_pir_api():
+            model_filename = self.pir_model_filename
+        else:
+            model_filename = self.model_filename
+
         exe = base.Executor(place)
         [
             inference_program,
@@ -509,7 +517,7 @@ def predict_static(self, data):
         ] = paddle.static.io.load_inference_model(
             self.model_save_dir,
             executor=exe,
-            model_filename=self.model_filename,
+            model_filename=model_filename,
             params_filename=self.params_filename,
         )
 
@@ -545,7 +553,7 @@ def verify_predict(self):
         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)
+
         np.testing.assert_allclose(
             dy_pre,
             st_pre,
@@ -559,16 +567,18 @@ def verify_predict(self):
             err_msg=f'dy_jit_pre:\n {dy_jit_pre}\n, st_pre: \n{st_pre}.',
         )
 
-        flat_st_pre = st_pre.flatten()
-        flat_predictor_pre = np.array(predictor_pre).flatten()
-        for i in range(len(flat_predictor_pre)):
-            # modify precision to 1e-6, avoid unittest failed
-            self.assertAlmostEqual(
-                flat_predictor_pre[i],
-                flat_st_pre[i],
-                delta=1e-6,
-                msg=f"predictor_pre:\n {flat_predictor_pre[i]}\n, st_pre: \n{flat_st_pre[i]}.",
-            )
+        if not use_pir_api():
+            predictor_pre = self.predict_analysis_inference(image)
+            flat_st_pre = st_pre.flatten()
+            flat_predictor_pre = np.array(predictor_pre).flatten()
+            for i in range(len(flat_predictor_pre)):
+                # modify precision to 1e-6, avoid unittest failed
+                self.assertAlmostEqual(
+                    flat_predictor_pre[i],
+                    flat_st_pre[i],
+                    delta=1e-6,
+                    msg=f"predictor_pre:\n {flat_predictor_pre[i]}\n, st_pre: \n{flat_st_pre[i]}.",
+                )
 
     @test_default_and_pir
     def test_check_result(self):
@@ -605,9 +615,7 @@ def test_check_result(self):
             err_msg=f'static acc5: {acc5_1} \ndygraph acc5: {acc5_2}',
         )
 
-        # TODO(@xiongkun): open after save / load supported in pir.
-        if not paddle.base.framework.use_pir_api():
-            self.verify_predict()
+        self.verify_predict()
 
 
 if __name__ == '__main__':
