From afafb1c347cb2e3331a21a68fa9df9b322c56712 Mon Sep 17 00:00:00 2001 From: 0x45f <23097963+0x45f@users.noreply.github.com> Date: Mon, 14 Mar 2022 10:12:27 +0800 Subject: [PATCH] Refine partial_program for new run_program OP (#40355) * refine partial_program * fix code for test_mnist.py train * support quantify UT * make __fake_vars and _double_grads to lazy * fix comments --- .../eager/to_static/run_program_op_func.h | 1 + .../eager/to_static/run_program_op_node.h | 23 +++++++++--- .../fluid/pybind/custom_handwrite_op_funcs.h | 1 - paddle/fluid/pybind/eager_method.cc | 17 +++++++++ .../dygraph_to_static/function_spec.py | 2 +- .../dygraph_to_static/partial_program.py | 35 ++++++++++--------- 6 files changed, 55 insertions(+), 24 deletions(-) diff --git a/paddle/fluid/eager/to_static/run_program_op_func.h b/paddle/fluid/eager/to_static/run_program_op_func.h index 6f8bccd64e4..9967d8c3690 100644 --- a/paddle/fluid/eager/to_static/run_program_op_func.h +++ b/paddle/fluid/eager/to_static/run_program_op_func.h @@ -57,6 +57,7 @@ inline void run_program_dygraph_function( auto grad_node = std::make_shared(1, 2); grad_node->SetFwdOutNames(out_names); + grad_node->SetOut(out); // Set Attributes grad_node->SetAttrMap(attrs); // Set TensorWrappers diff --git a/paddle/fluid/eager/to_static/run_program_op_node.h b/paddle/fluid/eager/to_static/run_program_op_node.h index ae5d86664a3..d99624e4932 100644 --- a/paddle/fluid/eager/to_static/run_program_op_node.h +++ b/paddle/fluid/eager/to_static/run_program_op_node.h @@ -260,9 +260,9 @@ inline void RunProgramAPI( } VLOG(2) << "The number of sub scopes after forward: " << out_scope_vec->front()->kids().size(); - // #ifdef PADDLE_WITH_MKLDNN - // if (FLAGS_use_mkldnn) paddle::platform::DontClearMKLDNNCache(place); - // #endif +#ifdef PADDLE_WITH_MKLDNN + if (FLAGS_use_mkldnn) paddle::platform::DontClearMKLDNNCache(place); +#endif } inline void RunProgramGradAPI( @@ -357,7 +357,7 @@ inline void RunProgramGradAPI( details::ShareTensorsFromScope(params_grad, *global_block, &scope); // Step5. drop current scope - // global_inner_scope->DeleteScope(&scope); + global_inner_scope->DeleteScope(&scope); VLOG(2) << "The number of sub scopes after backward: " << global_inner_scope->kids().size(); } @@ -400,6 +400,10 @@ class GradNodeRunProgram : public egr::GradNodeBase { paddle::platform::errors::InvalidArgument( "The grads[0].size() and fwd_out_names_.size() should be equal.")); for (size_t i = 0; i < fwd_out_names_.size(); ++i) { + auto &out_grad = egr::EagerUtils::unsafe_autograd_meta(*out_[i])->Grad(); + const_cast(out_grad).set_impl( + grads[0][i].impl()); + const_cast(grads[0][i]) .set_name(fwd_out_names_[i] + "@GRAD"); } @@ -432,6 +436,10 @@ class GradNodeRunProgram : public egr::GradNodeBase { fwd_out_names_ = out_names; } + void SetOut(const std::vector &out) { + out_ = out; + } + protected: void ConstructGradTensors( const std::vector &fwd_tensors, @@ -440,7 +448,11 @@ class GradNodeRunProgram : public egr::GradNodeBase { // such as: name, tensor type(DenseTensor or SelectedRows). VLOG(3) << "fwd_tensors.size(): " << fwd_tensors.size(); for (auto &fwd_t : fwd_tensors) { - grad_tensors->emplace_back(fwd_t.impl()); + if (phi::DenseTensor::classof(fwd_t.impl().get())) { + grad_tensors->emplace_back(std::make_shared()); + } else if (phi::SelectedRows::classof(fwd_t.impl().get())) { + grad_tensors->emplace_back(std::make_shared()); + } auto &grad_t = grad_tensors->back(); grad_t.set_name(fwd_t.name() + "@GRAD"); } @@ -462,6 +474,7 @@ class GradNodeRunProgram : public egr::GradNodeBase { std::vector step_scope_; std::vector fwd_out_names_; + std::vector out_; // Attribute Map paddle::framework::AttributeMap attrs_; diff --git a/paddle/fluid/pybind/custom_handwrite_op_funcs.h b/paddle/fluid/pybind/custom_handwrite_op_funcs.h index 7a276df0d5b..3b898ce77ce 100644 --- a/paddle/fluid/pybind/custom_handwrite_op_funcs.h +++ b/paddle/fluid/pybind/custom_handwrite_op_funcs.h @@ -31,7 +31,6 @@ static PyObject *eager_api_run_program(PyObject *self, PyObject *args, tstate = PyEval_SaveThread(); run_program_dygraph_function(X, Params, Out, OutScope, DOut, attrs); - std::cout << "end run_program_dygraph_function" << std::endl; PyEval_RestoreThread(tstate); tstate = nullptr; } catch (...) { diff --git a/paddle/fluid/pybind/eager_method.cc b/paddle/fluid/pybind/eager_method.cc index e5f22338dc6..b0dbce34d34 100644 --- a/paddle/fluid/pybind/eager_method.cc +++ b/paddle/fluid/pybind/eager_method.cc @@ -688,6 +688,21 @@ static PyObject* tensor_register_reduce_hook(TensorObject* self, PyObject* args, EAGER_CATCH_AND_THROW_RETURN_NULL } +static PyObject* set_grad_type(TensorObject* self, PyObject* args, + PyObject* kwargs) { + EAGER_TRY + auto var_type = pybind::CastPyArg2ProtoType(PyTuple_GET_ITEM(args, 0), 0); + auto grad_tensor = + egr::EagerUtils::unsafe_autograd_meta(self->tensor)->Grad(); + if (var_type == framework::proto::VarType::LOD_TENSOR) { + grad_tensor.set_impl(std::make_shared()); + } else if (var_type == framework::proto::VarType::SELECTED_ROWS) { + grad_tensor.set_impl(std::make_shared()); + } + return Py_None; + EAGER_CATCH_AND_THROW_RETURN_NULL +} + PyMethodDef variable_methods[] = { {"numpy", (PyCFunction)(void (*)(void))tensor_method_numpy, METH_VARARGS | METH_KEYWORDS, NULL}, @@ -734,6 +749,8 @@ PyMethodDef variable_methods[] = { {"_register_backward_hook", (PyCFunction)(void (*)(void))tensor_register_reduce_hook, METH_VARARGS | METH_KEYWORDS, NULL}, + {"_set_grad_type", (PyCFunction)(void (*)(void))set_grad_type, + METH_VARARGS | METH_KEYWORDS, NULL}, {NULL, NULL, 0, NULL}}; } // namespace pybind diff --git a/python/paddle/fluid/dygraph/dygraph_to_static/function_spec.py b/python/paddle/fluid/dygraph/dygraph_to_static/function_spec.py index 30012fb8666..900541459f6 100644 --- a/python/paddle/fluid/dygraph/dygraph_to_static/function_spec.py +++ b/python/paddle/fluid/dygraph/dygraph_to_static/function_spec.py @@ -104,7 +104,7 @@ class FunctionSpec(object): if isinstance(input_var, np.ndarray): input_var = paddle.static.InputSpec.from_numpy(input_var) _set_spec_stop_gradient(input_var, True) - elif isinstance(input_var, core.VarBase): + elif isinstance(input_var, (core.VarBase, core.eager.Tensor)): stop_gradient = input_var.stop_gradient input_var = paddle.static.InputSpec.from_tensor(input_var) _set_spec_stop_gradient(input_var, stop_gradient) diff --git a/python/paddle/fluid/dygraph/dygraph_to_static/partial_program.py b/python/paddle/fluid/dygraph/dygraph_to_static/partial_program.py index a442a8b92b6..216f955b751 100644 --- a/python/paddle/fluid/dygraph/dygraph_to_static/partial_program.py +++ b/python/paddle/fluid/dygraph/dygraph_to_static/partial_program.py @@ -148,10 +148,7 @@ class PartialProgramLayer: self._origin_main_program = self._verify_program(main_program) self._tmp_scope_vec = self._create_scope_vec() - # A fake_var to handle empty input or output - self.__fake_vars = _create_fake_var() # Set default mode to train - self._double_grads = self._get_double_grads(self._origin_main_program) self.training = True custom_white_list, custom_black_list = None, None @@ -163,6 +160,14 @@ class PartialProgramLayer: custom_white_list=custom_white_list, custom_black_list=custom_black_list) + @LazyInitialized + def __fake_vars(self): + return _create_fake_var() + + @LazyInitialized + def _double_grads(self): + return self._get_double_grads(self._origin_main_program) + @LazyInitialized def _infer_program(self): """ @@ -356,8 +361,10 @@ class PartialProgramLayer: def drop_scope_if_no_grad(self): tracer = framework._dygraph_tracer() + scope = self._tmp_scope_vec.value().get_scope() if isinstance( + self._tmp_scope_vec, (core.VarBase)) else self._tmp_scope_vec[0] if self.training and not tracer._has_grad: - self._tmp_scope_vec.value().get_scope().drop_kids() + scope.drop_kids() @property def program(self): @@ -449,18 +456,14 @@ class PartialProgramLayer: def _create_scope_vec(self): # Hold forward variables tmp_scope_vec = None + inner_scope = core.Scope() if not core._in_eager_mode(): tmp_scope_vec = core.VarBase(core.VarDesc.VarType.FP32, [], "program_out_scope", core.VarDesc.VarType.STEP_SCOPES, True) - # TODO(jiabin): Support this later. - # else: - # tmp_scope_vec = core.eager.Tensor(core.VarDesc.VarType.FP32, [], - # "program_out_scope", - # core.VarDesc.VarType.STEP_SCOPES, True) - - inner_scope = core.Scope() tmp_scope_vec.value().set_scope(inner_scope) + else: + tmp_scope_vec = [inner_scope] return tmp_scope_vec def _restore_out(self, out_vars): @@ -598,12 +601,10 @@ def _create_fake_var(): core.VarDesc.VarType.RAW, False) ] else: - return [] - # TODO(jiabin): Support this later - # return [ - # core.eager.Tensor(core.VarDesc.VarType.FP32, [], "Fake_var", - # core.VarDesc.VarType.RAW, False) - # ] + return [ + core.eager.Tensor(core.VarDesc.VarType.FP32, [], "Fake_var", + core.VarDesc.VarType.RAW, False) + ] def partial_program_from(concrete_program): -- GitLab