[pir_save_load] add jit load (train) (#63910)

xiaoguoguo626807 · web-flow · commit 007824fa4dfe · 2024-04-28T18:40:38.000+08:00
* add jit load.train

* modify backward program lost

* modify

* combine eval and train
diff --git a/paddle/fluid/pir/serialize_deserialize/include/interface.h b/paddle/fluid/pir/serialize_deserialize/include/interface.h
@@ -53,13 +53,14 @@ void WriteModule(const pir::Program& program,
  * deserilize program will be stored.
  * @param[in] pir_version  The current version of the PIR program format.
  *
- * @return Void. The function modifies the 'program' object to contain the data
- * read from the file.
+ * @return bool. The function modifies the 'program' object to contain the data
+ * read from the file. return bool indicates whether the program can use to
+ * funtune.
  *
  * @note If 'pir_version' is larger than the version of file, will trigger
  * version compatibility modification rule.
  */
-void ReadModule(const std::string& file_path,
+bool ReadModule(const std::string& file_path,
                 pir::Program* program,
                 const uint64_t& pir_version);
 
diff --git a/paddle/fluid/pir/serialize_deserialize/src/interface.cc b/paddle/fluid/pir/serialize_deserialize/src/interface.cc
@@ -23,6 +23,7 @@ namespace pir {
 #define BASE_CODE "base_code"
 #define MAGIC "magic"
 #define PIRVERSION "version"
+#define TRAINABLE "trainable"
 #define PIR "pir"
 void WriteModule(const pir::Program& program,
                  const std::string& file_path,
@@ -41,7 +42,8 @@ void WriteModule(const pir::Program& program,
   // write base code
   Json total;
 
-  total[BASE_CODE] = {{MAGIC, PIR}, {PIRVERSION, pir_version}};
+  total[BASE_CODE] = {
+      {MAGIC, PIR}, {PIRVERSION, pir_version}, {TRAINABLE, trainable}};
 
   ProgramWriter writer(pir_version, trainable);
   // write program
@@ -63,7 +65,7 @@ void WriteModule(const pir::Program& program,
   fout.close();
 }
 
-void ReadModule(const std::string& file_path,
+bool ReadModule(const std::string& file_path,
                 pir::Program* program,
                 const uint64_t& pir_version) {
   std::ifstream f(file_path);
@@ -83,6 +85,12 @@ void ReadModule(const std::string& file_path,
 
   ProgramReader reader(pir_version);
   reader.RecoverProgram(&(data[PROGRAM]), program);
+
+  if (data[BASE_CODE].contains(TRAINABLE)) {
+    return data[BASE_CODE][TRAINABLE].get<bool>();
+  } else {
+    return false;
+  }
 }
 
 }  // namespace pir
diff --git a/python/paddle/jit/pir_translated_layer.py b/python/paddle/jit/pir_translated_layer.py
@@ -17,7 +17,6 @@
 import numpy as np
 
 import paddle
-from paddle import _legacy_C_ops
 from paddle.base import core, framework, unique_name
 from paddle.base.dygraph.base import switch_to_static_graph
 from paddle.framework import in_dynamic_mode
@@ -36,9 +35,11 @@
 
 def _load_pir_program(model_file_path):
     program = paddle.static.Program()
-    paddle.base.core.deserialize_pir_program(model_file_path, program, 1)
+    trainable = paddle.base.core.deserialize_pir_program(
+        model_file_path, program, 1
+    )
 
-    return program
+    return program, trainable
 
 
 @switch_to_static_graph
@@ -58,7 +59,7 @@ def _get_pir_persistable_var_names(program):
 
 
 class _PirProgramHolder:
-    def __init__(self, program):
+    def __init__(self, program, trainable):
         super().__init__()
 
         # input, output, persistable,
@@ -67,10 +68,9 @@ def __init__(self, program):
         self._persistable_vars = []
         self._persistable_names = []
 
-        # execution scope
-        self._inner_scope = core.Scope()
-
+        self.support_train = trainable
         self._infer_program = program
+
         self._preprocess()
 
     def _preprocess(self):
@@ -124,14 +124,10 @@ def persistable_names(self):
     def persistable_vars(self):
         return self._persistable_vars
 
-    @property
-    def scope(self):
-        return self._inner_scope
-
 
-# [ PirTranslatedLayer : Run program in imperative mode ]
+# [ PirTranslatedLayer : Run program in dygraph mode ]
 #
-# DESIGN IDEA: using an special operator `RunProgram`, execute program inside operator.
+# DESIGN IDEA: using an special operator `PirRunProgram`, execute program inside operator.
 #
 # Op's Inputs:
 #   - the input variable of the user feed
@@ -228,9 +224,9 @@ def _construct_program_holders(model_path, model_filename=None):
                     continue
             else:
                 continue
-
+            program, trainable = _load_pir_program(model_file_path)
             program_holder_dict[func_name] = _PirProgramHolder(
-                _load_pir_program(model_file_path)
+                program, trainable
             )
     else:
         for _, _, file_names in os.walk(model_path):
@@ -242,8 +238,9 @@ def _construct_program_holders(model_path, model_filename=None):
                         method_name = 'forward'
                     else:
                         method_name.replace('model', '')
+                    program, trainable = _load_pir_program(model_file_path)
                     program_holder_dict[func_name] = _PirProgramHolder(
-                        _load_pir_program(model_file_path)
+                        program, trainable
                     )
 
     return program_holder_dict
@@ -264,10 +261,6 @@ def _construct_params_and_buffers(
         return var_dict
 
 
-def _valid_vars(vars):
-    return vars if vars else None
-
-
 def _run_dygraph(instance, input, program_holder):
     # 1. prepare inputs, outputs, attrs
     input_tensors = []
@@ -295,10 +288,6 @@ def _run_dygraph(instance, input, program_holder):
             tensor.name = program_holder.input_vars[i].name
         input_tensor_names.append(tensor.name)
         input_tensors.append(tensor)
-    if instance._input_args_names is None:
-        instance._input_args_names = [
-            ins.name for ins in program_holder.input_vars
-        ]
 
     persistable_tensors = []
     for var_name in program_holder.persistable_names:
@@ -313,72 +302,34 @@ def _run_dygraph(instance, input, program_holder):
                 % var_name
             )
 
-    output_tensors = []
+    from paddle.jit.dy2static.pir_partial_program import PartialProgramLayer
 
-    for i, var in enumerate(program_holder.output_vars):
-        tensor = core.eager.Tensor(
-            dtype=datatype_to_vartype[var.dtype],
-            dims=var.shape,
-            name=var.name,
-            type=core.VarDesc.VarType.LOD_TENSOR,
-            persistable=False,
-        )
-        output_tensors.append(tensor)
-
-    # hold forward variables
-    tmp_scope_vec = [program_holder.scope]
-
-    # 2. run program by op
-
-    forward_program = (
-        program_holder._infer_program
-        if instance._is_test
-        else program_holder.forward_program
-    )
+    inputs = program_holder.input_vars
+    outputs = program_holder.output_vars
+    parameters = (persistable_tensors, program_holder.persistable_vars)
 
-    backward_program = (
-        paddle.pir.Program()
-        if instance._is_test
-        else program_holder.backward_program
+    layer = PartialProgramLayer(
+        program_holder.infer_program,
+        inputs,
+        outputs,
+        parameters,
     )
+    instance.layer = layer
 
-    attrs = [
-        'forward_global_block',
-        forward_program.global_block(),
-        'backward_global_block',
-        backward_program.global_block(),
-        'is_test',
-        instance._is_test,
-        'program_id',
-        paddle.utils._hash_with_id(forward_program, instance),
-        'fx',
-        program_holder.input_vars,
-        'fo',
-        program_holder.output_vars,
-        'fp',
-        program_holder.persistable_vars,
-        'fm',
-        [],
-        'no_need_buffers',
-        [],
-    ]
-
-    _legacy_C_ops.pir_run_program(
-        _valid_vars(input_tensors),
-        _valid_vars(persistable_tensors),
-        _valid_vars(output_tensors),
-        tmp_scope_vec,
-        None,
-        *attrs,
-    )
+    if instance._is_test:
+        layer.training = False
+    else:
+        if not program_holder.support_train:
+            raise ValueError(
+                "The model is not trainable, please check model_file of jit.save."
+            )
+        else:
+            layer.training = True
 
-    outs = output_tensors
-    if len(output_tensors) == 1:
-        outs = output_tensors[0]
-    return outs
+    return instance.layer(input_tensors)
 
 
-def _run_static_graph(input, program_holder, trace_program):
+def _run_static_graph(program_holder, trace_program):
     paddle.base.framework.switch_main_program(trace_program)
     return program_holder.output_vars
 
@@ -594,7 +545,7 @@ def __i_m_p_l__(self, *input):
                 return _run_dygraph(self, input, program_holder)
             else:
                 return _run_static_graph(
-                    input, program_holder, program_holder.infer_program
+                    program_holder, program_holder.infer_program
                 )
 
         __i_m_p_l__.__name__ = method_name
diff --git a/test/deprecated/legacy_test/test_io_save_load.py b/test/deprecated/legacy_test/test_io_save_load.py
@@ -101,18 +101,35 @@ def tearDown(self):
     def test_when_train_with_no_grad(self):
         paddle.disable_static()
         net = paddle.nn.Linear(1024, 1)
-        net = paddle.jit.to_static(net)
+        net = paddle.jit.to_static(net, full_graph=True)
         x = paddle.rand([1024], 'float32')
+        x.stop_gradient = False
         out = net(x)
-        save_path = os.path.join(self.temp_dir.name, 'train_with_no_grad')
+        out.backward()
+        x_grad = x.grad.mean()
+        x.clear_grad()
 
+        # jit.save
+        save_path = os.path.join(self.temp_dir.name, 'train_with_no_grad')
         paddle.jit.save(net, save_path)
-        net = paddle.jit.load(save_path)
-        net.eval()
+
+        # test eval mode
+        net1 = paddle.jit.load(save_path)
+        net1.eval()
 
         with paddle.no_grad():
-            out2 = net(x)
-            self.assertEqual(out, out2)
+            out1 = net1(x)
+            self.assertEqual(out, out1)
+
+        # test train mode
+        net2 = paddle.jit.load(save_path)
+        net2.train()
+        out2 = net2(x)
+        out2.backward()
+        self.assertEqual(out, out2)
+        x_grad2 = x.grad.mean()
+        if paddle.framework.in_pir_mode():
+            self.assertEqual(x_grad, x_grad2)
 
 
 if __name__ == '__main__':
