diff --git a/python/paddle/jit/dy2static/partial_program.py b/python/paddle/jit/dy2static/partial_program.py index a1fc67d04e6c60b975a4be6c9122fb5051412f00..0264d645a0676b41aba816cf4dcf3c4a6a56115d 100644 --- a/python/paddle/jit/dy2static/partial_program.py +++ b/python/paddle/jit/dy2static/partial_program.py @@ -33,6 +33,7 @@ from .utils import ( RETURN_NO_VALUE_MAGIC_NUM, backend_guard, construct_grad_names, + tensor_name_guard, ) __all__ = [] @@ -223,30 +224,25 @@ class PartialProgramLayer: Execute static graph by Interpreter and Return dynamic Tensors. """ with UniqueNameGuard(self._name_generator): - in_vars, out_vars, in_var_names, resume_name_record = self._prepare( - inputs - ) + in_vars, out_vars, in_var_names = 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] + with tensor_name_guard(in_vars, in_var_names): + _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) @@ -905,7 +901,6 @@ class PartialProgramLayer: # 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): @@ -928,8 +923,6 @@ class PartialProgramLayer: 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()) @@ -960,7 +953,7 @@ class PartialProgramLayer: # Create Tensor to receive output data. out_vars = list(map(create_out, self._outputs.var_ids)) - return input_vars, out_vars, input_var_names, resume_name_record + return input_vars, out_vars, input_var_names def _create_scope_vec(self, program_id=None, use_scope_cache=False): # Hold forward variables diff --git a/python/paddle/jit/dy2static/utils.py b/python/paddle/jit/dy2static/utils.py index eaf808fbf4f075f7cc321d08825e9319351f1637..eaa47258b10742634126f69e6db52e693998c9e1 100644 --- a/python/paddle/jit/dy2static/utils.py +++ b/python/paddle/jit/dy2static/utils.py @@ -1527,3 +1527,16 @@ def construct_grad_names(grad_info_map, x_vars, param_vars, out_vars): out_grad_vars = backward._get_grad_vars(grad_info_map, out_vars) grad_var_names['out'] = list(map(fn, out_grad_vars)) return grad_var_names + + +@signature_safe_contextmanager +def tensor_name_guard(tensors, names): + try: + assert len(tensors) == len(names) + origin_names = [t.name for t in tensors] + for t, name in zip(tensors, names): + t.name = name + yield + finally: + for t, name in zip(tensors, origin_names): + t.name = name