【pir save load 】modify Accuracy diff of Test bert (#64168)

xiaoguoguo626807 · web-flow · commit e938f0eb2e1e · 2024-05-11T11:47:39.000+08:00
diff --git a/paddle/fluid/pybind/pir.cc b/paddle/fluid/pybind/pir.cc
@@ -666,6 +666,12 @@ void BindOperation(py::module *m) {
           [](Operation &self) { return &self.blocks(); },
           return_value_policy::reference)
       .def("has_attr", &Operation::HasAttribute)
+      .def("set_bool_attr",
+           [](Operation &self, std::string &attr_name, bool flag) {
+             self.set_attribute(
+                 attr_name,
+                 pir::BoolAttribute::get(pir::IrContext::Instance(), flag));
+           })
       .def("attrs",
            [](Operation &self) -> py::dict {
              py::dict attrs_dict;
diff --git a/python/paddle/static/pir_io.py b/python/paddle/static/pir_io.py
@@ -159,6 +159,42 @@ def pir_prune_with_input(program, feed_vars, target_vars):
             program.global_block().remove_op(op)
 
 
+def _inference_optimize(program, prune_read_op=True):
+    """
+    This method will create a new program and do following adjustments on it:
+    1. Remove all reader variables and their creator ops if exist.
+
+    2. Remove the :code:`read_op` if exists.
+
+    3. change the :code:`is_test`
+    attribute of operators to :code:`True`. All the :code:`Parameter`
+    information will be lost.
+
+    Args:
+        prune_read_op(bool): remove the read ops that are added by py_reader
+                             for cpp inference library
+
+    Notes: This API is a very low level API. Use
+    :code:`Program.clone(for_test=True)` instead.
+
+    Returns:
+        Program: The new program.
+    """
+
+    # 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 ")
+
+    # change all `is_test` attributes to True
+    for block in program.blocks:
+        for op in block.ops:
+            if op.has_attr("is_test"):
+                op.set_bool_attr("is_test", True)
+            if op.name() == "pd_op.batch_norm":
+                # Remove the output ReserveSpace of batch_norm if exists.
+                pass
+
+
 def normalize_pir_program(program, feed_vars, fetch_vars, **kwargs):
     """
 
@@ -242,7 +278,7 @@ def normalize_pir_program(program, feed_vars, fetch_vars, **kwargs):
     # if feed var is not conect with target_vars, it will be delete.
     if not skip_prune_program:
         pir_prune_with_input(copy_program, clone_feed_vars, clone_fetch_vars)
-    # copy_program = copy_program._inference_optimize(prune_read_op=True)
+    _inference_optimize(copy_program, prune_read_op=True)
 
     fetch_vars_tuple = []
     for i, var in enumerate(clone_fetch_vars):
@@ -415,6 +451,11 @@ def load_vars_pir(
         if main_program is None:
             main_program = default_main_program()
 
+        if not isinstance(main_program, paddle.static.Program):
+            raise TypeError(
+                "The type of input main_program is invalid, expected type is paddle.static.Program, but received %s"
+                % type(main_program)
+            )
         param, opt = get_pir_parameters(main_program)
         vars_list = param + opt
         load_vars_pir(
diff --git a/test/dygraph_to_static/test_bert.py b/test/dygraph_to_static/test_bert.py
@@ -320,13 +320,13 @@ def verify_predict(self):
                     np.testing.assert_allclose(
                         st_res,
                         dy_res,
-                        rtol=1e-04,
+                        rtol=1e-05,
                         err_msg=f'dygraph_res: {dy_res[~np.isclose(st_res, dy_res)]},\n static_res: {st_res[~np.isclose(st_res, dy_res)]}',
                     )
                     np.testing.assert_allclose(
                         st_res,
                         dy_jit_res,
-                        rtol=1e-04,
+                        rtol=1e-05,
                         err_msg=f'dygraph_jit_res: {dy_jit_res[~np.isclose(st_res, dy_jit_res)]},\n static_res: {st_res[~np.isclose(st_res, dy_jit_res)]}',
                     )
             else:
