PaddlePaddle
diff --git a/‎paddle/fluid/eager/to_static/run_program_op_func.h
Copy file name to clipboard
+7-3Lines changed: 7 additions & 3 deletions b/‎paddle/fluid/eager/to_static/run_program_op_func.h
Copy file name to clipboard
+7-3Lines changed: 7 additions & 3 deletions
diff --git a/‎paddle/fluid/eager/to_static/run_program_op_node.h
Copy file name to clipboard
+29-3Lines changed: 29 additions & 3 deletions b/‎paddle/fluid/eager/to_static/run_program_op_node.h
Copy file name to clipboard
+29-3Lines changed: 29 additions & 3 deletions
diff --git a/‎paddle/fluid/operators/run_program_op.cc
Copy file name to clipboard
+4Lines changed: 4 additions & 0 deletions b/‎paddle/fluid/operators/run_program_op.cc
Copy file name to clipboard
+4Lines changed: 4 additions & 0 deletions
diff --git a/‎python/paddle/jit/api.py
Copy file name to clipboard
+22-5Lines changed: 22 additions & 5 deletions b/‎python/paddle/jit/api.py
Copy file name to clipboard
+22-5Lines changed: 22 additions & 5 deletions
diff --git a/‎python/paddle/jit/dy2static/function_spec.py
Copy file name to clipboard
+49-2Lines changed: 49 additions & 2 deletions b/‎python/paddle/jit/dy2static/function_spec.py
Copy file name to clipboard
+49-2Lines changed: 49 additions & 2 deletions
diff --git a/‎python/paddle/jit/dy2static/partial_program.py
Copy file name to clipboard
+44-22Lines changed: 44 additions & 22 deletions b/‎python/paddle/jit/dy2static/partial_program.py
Copy file name to clipboard
+44-22Lines changed: 44 additions & 22 deletions
@@ -77,10 +77,11 @@ static void clear_unused_out_var_in_backward(
 
 static std::vector<paddle::Tensor> filter_unused_input_var_in_backward(
     const std::vector<paddle::Tensor>& x,
+    const std::vector<std::string>& x_names,
     const paddle::framework::BlockDesc* backward_block) {
   auto filter_x = std::vector<paddle::Tensor>(x);
   for (size_t i = 0; i < x.size(); i++) {
-    if (!backward_block->HasVar(x[i].name())) {
+    if (!backward_block->HasVar(x_names[i])) {
       auto fake = paddle::Tensor(std::make_shared<phi::DenseTensor>());
       fake.set_name(paddle::framework::kFakeVarName);
       filter_x[i] = fake;
@@ -117,6 +118,9 @@ inline void run_program_ad_func(
   VLOG(2) << "start run run_program grad";
 
   if (require_any_grad) {
+    auto x_names =
+        PADDLE_GET_CONST(std::vector<std::string>, attrs.at("x_names"));
+
     egr::EagerUtils::PassStopGradient(false, &p_autograd_outs);
     // Create GradOpNode (1 means [out_grad], 2 means [x_grad, paramx_grad])
     auto grad_node = std::make_shared<GradNodeRunProgram>(1, 2);
@@ -130,7 +134,7 @@ inline void run_program_ad_func(
         paddle::framework::BlockDesc*, attrs.at("backward_global_block"));
     // Clear unused x vars
     auto filter_x =
-        filter_unused_input_var_in_backward(x, backward_global_block);
+        filter_unused_input_var_in_backward(x, x_names, backward_global_block);
     // Set TensorWrappers
     grad_node->SetFwdX(filter_x);
     // Clear unused out vars
@@ -145,7 +149,7 @@ inline void run_program_ad_func(
 
     std::vector<const paddle::Tensor*> x_require_grad;
     for (size_t i = 0; i < x.size(); ++i) {
-      auto& name = x[i].name();
+      auto& name = x_names[i];
       if (forward_global_block->HasVar(name) ||
           backward_global_block->HasVar(name)) {
         x_require_grad.push_back(&x[i]);
 
@@ -150,6 +150,31 @@ static void ShareTensorsIntoScope(const std::vector<Tensor> &tensors,
   }
 }
 
+static void ShareTensorsIntoScopeWithName(
+    const std::vector<Tensor> &tensors,
+    const std::vector<std::string> &tensor_names,
+    paddle::framework::Scope *scope) {
+  for (size_t i = 0; i < tensors.size(); ++i) {
+    auto name = tensor_names[i];
+    if (name == paddle::framework::kFakeVarName) {
+      continue;
+    }
+    auto *var = scope->Var(name);
+    CheckInputVarStatus(tensors[i]);
+    // share tensor
+    auto tensor_base = tensors[i].impl();
+    if (phi::DenseTensor::classof(tensor_base.get())) {
+      auto *dst_tensor = var->GetMutable<phi::DenseTensor>();
+      auto t = std::dynamic_pointer_cast<phi::DenseTensor>(tensor_base);
+      *dst_tensor = *t;
+    } else if (phi::SelectedRows::classof(tensor_base.get())) {
+      auto *dst_tensor = var->GetMutable<phi::SelectedRows>();
+      auto t = std::dynamic_pointer_cast<phi::SelectedRows>(tensor_base);
+      *dst_tensor = *t;
+    }
+  }
+}
+
 static void ShareTensorsFromScope(
     const std::vector<Tensor *> &tensors,
     const paddle::framework::BlockDesc &global_block,
@@ -320,7 +345,8 @@ inline void RunProgramAPI(
 
   VLOG(4) << "global_inner_scope:" << global_inner_scope;
 
-  auto input_names = details::GetTensorsName(x);
+  auto input_names =
+      PADDLE_GET_CONST(std::vector<std::string>, attrs.at("x_names"));
   auto output_names = details::GetTensorsName(out);
   auto param_names = details::GetTensorsName(params);
   auto dout_names = details::GetTensorsName(dout);
@@ -371,7 +397,7 @@ inline void RunProgramAPI(
                "for program: "
             << program_id;
     // Step 1. share input_vars & parameters into scope
-    details::ShareTensorsIntoScope(x, global_inner_scope);
+    details::ShareTensorsIntoScopeWithName(x, input_names, global_inner_scope);
     details::ShareTensorsIntoScope(params, global_inner_scope);
     // Step 2. create new interpretercore
 
@@ -433,7 +459,7 @@ inline void RunProgramAPI(
         program_id, global_inner_scope, /*is_grad=*/false);
     interpreter_core = cached_value.core_;
     // Step 2. update scope for cache interpretercore
-    details::ShareTensorsIntoScope(x, global_inner_scope);
+    details::ShareTensorsIntoScopeWithName(x, input_names, global_inner_scope);
     details::ShareTensorsIntoScope(params, global_inner_scope);
     if (interpreter_core->GetVariableScope()->GetMutableScope() !=
         global_inner_scope) {
 
@@ -139,6 +139,10 @@ class RunProgramOpMaker : public framework::OpProtoAndCheckerMaker {
                                       "std::vector<std::string>"
                                       "The names of output gradients.")
         .SetDefault({});
+    AddAttr<std::vector<std::string>>("x_names",
+                                      "std::vector<std::string>"
+                                      "The names of input tensors.")
+        .SetDefault({});
     AddAttr<std::vector<std::string>>("x_grad_names",
                                       "std::vector<std::string>"
                                       "The names of input gradients.")
 
@@ -415,6 +415,9 @@ def __init__(self):
         # if True, multi `StaticFunction` will share params in one file.
         self.combine_params = False
 
+        # when need to save a prune model, use input_names_after_prune to specify the inputs left after pruning
+        self.input_names_after_prune = None
+
     @property
     def output_spec(self):
         return self._output_spec
@@ -488,11 +491,12 @@ def keep_name_table(self, value):
 
 def _parse_save_configs(configs):
     supported_configs = [
-        'output_spec',
+        "output_spec",
         "with_hook",
         "combine_params",
         "clip_extra",
         "skip_forward",
+        "input_names_after_prune",
     ]
 
     # input check
@@ -505,11 +509,14 @@ def _parse_save_configs(configs):
 
     # construct inner config
     inner_config = _SaveLoadConfig()
-    inner_config.output_spec = configs.get('output_spec', None)
-    inner_config.with_hook = configs.get('with_hook', False)
+    inner_config.output_spec = configs.get("output_spec", None)
+    inner_config.with_hook = configs.get("with_hook", False)
     inner_config.combine_params = configs.get("combine_params", False)
     inner_config.clip_extra = configs.get("clip_extra", True)
     inner_config.skip_forward = configs.get("skip_forward", False)
+    inner_config.input_names_after_prune = configs.get(
+        "input_names_after_prune", None
+    )
 
     return inner_config
 
@@ -533,7 +540,7 @@ def _parse_load_config(configs):
     return inner_config
 
 
-def _get_input_var_names(inputs, input_spec):
+def _get_input_var_names(inputs, input_spec, input_names_after_prune):
     name_none_error = (
         "The %s's name is None. "
         "When using jit.save, please set InputSepc's name in "
@@ -546,6 +553,14 @@ def _get_input_var_names(inputs, input_spec):
         "in input_spec is the same as the name of InputSpec in "
         "`to_static` decorated on the Layer.forward method."
     )
+    if input_names_after_prune is not None:
+        input_spec = [
+            x
+            for x in input_spec
+            if isinstance(x, paddle.static.InputSpec)
+            and x.name in input_names_after_prune
+        ]
+
     result_list = []
     input_var_names = [
         var.name
@@ -1201,7 +1216,9 @@ def save(layer, path, input_spec=None, **configs):
         #   - the input_spec length < len((concrete_program.inputs) - 1
         #   - the input_spec's name should be in concrete_program.inputs
         input_var_names = _get_input_var_names(
-            concrete_program.inputs, inner_input_spec
+            concrete_program.inputs,
+            inner_input_spec,
+            configs.input_names_after_prune,
         )
 
         # NOTE(chenweihang): [ Get output variables ]
 
@@ -150,8 +150,17 @@ def args_to_input_spec(self, args, kwargs):
             # replace argument with corresponding InputSpec.
             args_with_spec = convert_to_input_spec(args, self._input_spec)
         else:
-            args_with_spec = _replace_value_with_input_spec(args)
-            kwargs_with_spec = _replace_value_with_input_spec(kwargs)
+            args_with_spec = _replace_to_input_spec_with_new_name(
+                args, self._arg_names
+            )
+            kwarg_names = ["kwargs." + key for key in kwargs.keys()]
+            kwargs_list_with_spec = _replace_to_input_spec_with_new_name(
+                list(kwargs.values()), kwarg_names
+            )
+            kwargs_with_spec = {
+                key: kwargs_list_with_spec[idx]
+                for idx, key in enumerate(kwargs)
+            }
 
         # If without specificing name in input_spec, add default name
         # according to argument name from decorated function.
@@ -302,6 +311,44 @@ def _replace_value_with_input_spec(args):
     return args_with_spec
 
 
+def _replace_to_input_spec_with_new_name(args, arg_names):
+    assert len(args) == len(arg_names)
+    order_digit = len(str(len(arg_names) - 1))
+    args_with_spec = []
+    for order, (arg, name_prefix) in enumerate(zip(args, arg_names)):
+        index = 0
+        for idx, origin_input in enumerate(paddle.utils.flatten(arg)):
+            if isinstance(origin_input, np.ndarray):
+                input_var = paddle.static.InputSpec.from_numpy(origin_input)
+                input_var.stop_gradient = True
+            elif isinstance(origin_input, core.eager.Tensor):
+                stop_gradient = origin_input.stop_gradient
+                input_var = paddle.static.InputSpec.from_tensor(origin_input)
+                input_var.stop_gradient = stop_gradient
+            elif isinstance(origin_input, paddle.fluid.framework.Variable):
+                stop_gradient = origin_input.stop_gradient
+                input_var = paddle.static.InputSpec(
+                    origin_input.shape, origin_input.dtype, origin_input.name
+                )
+                input_var.stop_gradient = stop_gradient
+            else:
+                input_var = origin_input
+
+            if isinstance(
+                origin_input,
+                (
+                    np.ndarray,
+                    core.eager.Tensor,
+                    paddle.fluid.framework.Variable,
+                ),
+            ):
+                input_var.name = f"_jst.{str(order).zfill(order_digit)}.{name_prefix}.{str(index)}"
+                index += 1
+            args_with_spec.append(input_var)
+    args_with_spec = paddle.utils.pack_sequence_as(args, args_with_spec)
+    return args_with_spec
+
+
 def convert_to_input_spec(inputs, input_spec):
     """
     Replaces tensor in structured `inputs` by InputSpec in `input_spec`.
 
@@ -25,6 +25,7 @@
 from paddle.fluid.data_feeder import check_type, convert_dtype
 from paddle.fluid.dygraph.base import switch_to_static_graph
 from paddle.fluid.framework import _apply_pass
+from paddle.fluid.unique_name import guard as UniqueNameGuard
 from paddle.optimizer.lr import LRScheduler
 
 from . import logging_utils
@@ -170,12 +171,19 @@ class PartialProgramLayer:
     """
 
     def __init__(
-        self, main_program, inputs, outputs, parameters=None, **kwargs
+        self,
+        main_program,
+        inputs,
+        outputs,
+        name_generator,
+        parameters=None,
+        **kwargs
     ):
         super().__init__()
         self._inputs = NestSequence(inputs)
         self._outputs = NestSequence(outputs, need_check=True)
         self._params = parameters if parameters is not None else []
+        self._name_generator = name_generator
 
         self._build_strategy = kwargs.get('build_strategy', BuildStrategy())
         assert isinstance(self._build_strategy, BuildStrategy)
@@ -214,26 +222,35 @@ def __call__(self, inputs):
         """
         Execute static graph by Interpreter and Return dynamic Tensors.
         """
-        in_vars, out_vars = self._prepare(inputs)
-        self._cast_fp16_if_pure_fp16(in_vars)
-        attrs = self._prepare_attributes()
-
-        self._sync_lr_value_with_scheduler()
-
-        _legacy_C_ops.run_program(
-            self._valid_vars(in_vars),
-            self._valid_vars(self._params),
-            self._valid_vars(out_vars),
-            self._create_scope_vec(
-                program_id=self.program_id, use_scope_cache=True
-            ),
-            self._double_grads,
-            self._cuda_graph_vec,
-            *attrs
-        )
-        self._update_stop_gradient(out_vars)
-        restored_nest_out = self._restore_out(out_vars)
-        return self._remove_no_value(restored_nest_out)
+        with UniqueNameGuard(self._name_generator):
+            in_vars, out_vars, in_var_names, resume_name_record = self._prepare(
+                inputs
+            )
+            self._cast_fp16_if_pure_fp16(in_vars)
+            attrs = self._prepare_attributes()
+            attrs.extend(["x_names", in_var_names])
+
+            self._sync_lr_value_with_scheduler()
+
+            _legacy_C_ops.run_program(
+                self._valid_vars(in_vars),
+                self._valid_vars(self._params),
+                self._valid_vars(out_vars),
+                self._create_scope_vec(
+                    program_id=self.program_id, use_scope_cache=True
+                ),
+                self._double_grads,
+                self._cuda_graph_vec,
+                *attrs
+            )
+
+            for var in in_vars:
+                if var.name in resume_name_record:
+                    var.name = resume_name_record[var.name]
+
+            self._update_stop_gradient(out_vars)
+            restored_nest_out = self._restore_out(out_vars)
+            return self._remove_no_value(restored_nest_out)
 
     def _sync_lr_value_with_scheduler(self):
         """Update lr_var value with calculated by lr_scheduler."""
@@ -887,6 +904,8 @@ def _prepare(self, inputs):
         flatten_inputs = paddle.utils.flatten(inputs)
         # Convert variable into Tensor and feed in training data.
         input_vars = []
+        input_var_names = []
+        resume_name_record = {}
         expected_place = framework._current_expected_place()
         for i, value in enumerate(flatten_inputs):
             if isinstance(value, np.ndarray):
@@ -909,9 +928,11 @@ def _prepare(self, inputs):
                     var.stop_gradient = True
                 else:
                     var = value
+                resume_name_record[self._inputs[i].desc.name()] = var.name
                 var.name = self._inputs[i].desc.name()
             else:
                 continue
+            input_var_names.append(self._inputs[i].desc.name())
             input_vars.append(var)
 
         # mapping from name(string) -> Tensor
@@ -939,7 +960,7 @@ def create_out(var_id):
         # Create Tensor to receive output data.
         out_vars = list(map(create_out, self._outputs.var_ids))
 
-        return input_vars, out_vars
+        return input_vars, out_vars, input_var_names, resume_name_record
 
     def _create_scope_vec(self, program_id=None, use_scope_cache=False):
         # Hold forward variables
@@ -1106,6 +1127,7 @@ def partial_program_from(concrete_program, from_method=False):
         concrete_program.main_program,
         inputs,
         concrete_program.outputs,
+        concrete_program.name_generator,
         concrete_program.parameters,
         **concrete_program.kwargs
     )