From b3ff6043ae29436c37ea4300978add3d96a9e682 Mon Sep 17 00:00:00 2001 From: "7347157+joylvliang@user.noreply.gitee.com" Date: Tue, 11 Jan 2022 14:25:30 +0800 Subject: [PATCH] Implement_forward_hook_for_PyNative --- .../ccsrc/backend/graph_compiler/backend.cc | 74 ++-- .../ccsrc/frontend/optimizer/ad/dfunctor.cc | 75 ++-- .../ccsrc/frontend/optimizer/ad/dfunctor.h | 2 + .../ccsrc/frontend/optimizer/ad/kprim.cc | 46 ++- .../ccsrc/frontend/optimizer/ad/kpynative.cc | 7 +- .../optimizer/ad/prim_bprop_optimizer.cc | 3 +- mindspore/ccsrc/frontend/optimizer/irpass.cc | 8 +- .../optimizer/irpass/special_op_eliminate.h | 8 +- .../pipeline/jit/parse/data_converter.cc | 2 +- .../pipeline/pynative/pynative_execute.cc | 46 ++- .../pipeline/pynative/pynative_execute.h | 5 + mindspore/ccsrc/pybind_api/ir/primitive_py.cc | 298 +++++++++------ mindspore/ccsrc/pybind_api/ir/primitive_py.h | 43 +-- mindspore/core/base/core_ops.h | 1 + .../python/mindspore/_extends/parse/parser.py | 4 +- mindspore/python/mindspore/common/api.py | 3 + .../python/mindspore/common/hook_handle.py | 110 ++++++ mindspore/python/mindspore/nn/cell.py | 338 ++++++++++++++++-- .../mindspore/ops/operations/_inner_ops.py | 118 +++++- .../mindspore/ops/operations/debug_ops.py | 45 ++- 20 files changed, 965 insertions(+), 271 deletions(-) create mode 100644 mindspore/python/mindspore/common/hook_handle.py diff --git a/mindspore/ccsrc/backend/graph_compiler/backend.cc b/mindspore/ccsrc/backend/graph_compiler/backend.cc index b054a1a6e80..8cfa4c90654 100644 --- a/mindspore/ccsrc/backend/graph_compiler/backend.cc +++ b/mindspore/ccsrc/backend/graph_compiler/backend.cc @@ -580,6 +580,30 @@ void MindRTBackend::CompileGraph(const GraphSegmentPtr &segment) { } namespace { +ValuePtr GetControlOpInputFromMakeTuple(const std::shared_ptr &graph_compiler, + const AnfNodePtr &front_cnode, const CNodePtr &backend_cnode, + const std::map &op_output_map, + const std::map ¶meter_index, + const std::vector &graph_inputs, + InputTensorInfo *input_tensor_info, size_t *input_index) { + MS_EXCEPTION_IF_NULL(graph_compiler); + MS_EXCEPTION_IF_NULL(front_cnode); + MS_EXCEPTION_IF_NULL(input_index); + MS_LOG(DEBUG) << "The input node of hook op: " << front_cnode->DebugString() << " is a make tuple node."; + auto make_tuple = front_cnode->cast(); + MS_EXCEPTION_IF_NULL(make_tuple); + const auto output_size = make_tuple->size() - 1; + std::vector output_values; + for (size_t idx = 0; idx < output_size; ++idx) { + TensorPtr tensor = graph_compiler->GetSingleOpInputTensorByIndex(backend_cnode, op_output_map, parameter_index, + graph_inputs, input_tensor_info, *input_index); + MS_EXCEPTION_IF_NULL(tensor); + output_values.emplace_back(tensor); + ++(*input_index); + } + return std::make_shared(output_values); +} + void GetControlOpInput(const std::shared_ptr &graph_compiler, const CNodePtr &front_cnode, const CNodePtr &backend_cnode, const std::map &op_output_map, const std::map ¶meter_index, @@ -594,34 +618,36 @@ void GetControlOpInput(const std::shared_ptr &graph_compiler, con for (size_t i = 1; i < inputs.size(); i++) { const auto &input_node = inputs[i]; MS_EXCEPTION_IF_NULL(input_node); - auto kernel_with_index = AnfAlgo::VisitKernel(input_node, 0); - auto real_input = kernel_with_index.first; - MS_EXCEPTION_IF_NULL(real_input); - - if (!real_input->isa()) { - TensorPtr tensor = graph_compiler->GetSingleOpInputTensorByIndex(backend_cnode, op_output_map, parameter_index, - graph_inputs, input_tensor_info, input_index); - MS_EXCEPTION_IF_NULL(tensor); - args->emplace_back(tensor); - input_index++; + if (IsPrimitiveCNode(input_node, prim::kPrimMakeTuple)) { + // Hook multi-input or multi-output. + args->emplace_back(GetControlOpInputFromMakeTuple(graph_compiler, input_node, backend_cnode, op_output_map, + parameter_index, graph_inputs, input_tensor_info, + &input_index)); continue; } - - // Get value from value node. - const auto &value_node = real_input->cast(); - MS_EXCEPTION_IF_NULL(value_node); - const auto &value = value_node->value(); - MS_EXCEPTION_IF_NULL(value); - - if (value->isa()) { - const auto &value_sequeue = value->cast(); - MS_EXCEPTION_IF_NULL(value_sequeue); - input_index += value_sequeue->size(); + // Hook single-input or single-output. + auto real_input = AnfAlgo::VisitKernel(input_node, 0).first; + MS_EXCEPTION_IF_NULL(real_input); + if (!real_input->isa()) { + auto tensor = graph_compiler->GetSingleOpInputTensorByIndex(backend_cnode, op_output_map, parameter_index, + graph_inputs, input_tensor_info, input_index); + MS_EXCEPTION_IF_NULL(tensor); + args->emplace_back(tensor); + ++input_index; } else { - input_index++; + const auto &value_node = real_input->cast(); + MS_EXCEPTION_IF_NULL(value_node); + const auto &value = value_node->value(); + MS_EXCEPTION_IF_NULL(value); + args->emplace_back(value); + if (value->isa()) { + const auto &value_sequeue = value->cast(); + MS_EXCEPTION_IF_NULL(value_sequeue); + input_index += value_sequeue->size(); + } else { + ++input_index; + } } - - args->emplace_back(value); } } diff --git a/mindspore/ccsrc/frontend/optimizer/ad/dfunctor.cc b/mindspore/ccsrc/frontend/optimizer/ad/dfunctor.cc index 5eff538a86c..490af498e0a 100644 --- a/mindspore/ccsrc/frontend/optimizer/ad/dfunctor.cc +++ b/mindspore/ccsrc/frontend/optimizer/ad/dfunctor.cc @@ -186,6 +186,52 @@ static bool HasSideEffectBackProp(const CNodePtr &cnode) { return false; } +static AnfNodePtr SkipHookNodeInBackProp(const AnfNodePtr &node) { + MS_EXCEPTION_IF_NULL(node); + if (IsPrimitiveCNode(node, prim::kPrimHookBackward) || IsPrimitiveCNode(node, prim::kPrimCellBackwardHook)) { + MS_LOG(WARNING) + << "Hook operation does not work in graph mode or ms_function, it will be eliminated during compilation."; + auto output_cnode = node->cast(); + if (output_cnode->size() - 1 == 1) { + return output_cnode->input(1); + } + // Replace hook node with make tuple node. + abstract::AbstractBasePtrList multi_output_abs; + std::vector multi_output_nodes{NewValueNode(prim::kPrimMakeTuple)}; + std::for_each(output_cnode->inputs().begin() + 1, output_cnode->inputs().end(), + [&multi_output_nodes, &multi_output_abs](const AnfNodePtr &inp) { + MS_EXCEPTION_IF_NULL(inp); + multi_output_nodes.emplace_back(inp); + multi_output_abs.emplace_back(inp->abstract()); + }); + auto primal_graph = node->func_graph(); + MS_EXCEPTION_IF_NULL(primal_graph); + auto make_tuple = primal_graph->NewCNode(std::move(multi_output_nodes)); + make_tuple->set_abstract(std::make_shared(multi_output_abs)); + auto mng = primal_graph->manager(); + MS_EXCEPTION_IF_NULL(mng); + if (!mng->Replace(node, make_tuple)) { + MS_LOG(EXCEPTION) << "Failed to replace old node: " << node->DebugString() + << " with new node: " << make_tuple->DebugString(); + } + return make_tuple; + } + if (IsPrimitiveCNode(node, prim::kPrimTupleGetItem)) { + auto tuple_get_item = node->cast(); + auto inp = tuple_get_item->input(1); + if (IsPrimitiveCNode(inp, prim::kPrimHookBackward) || IsPrimitiveCNode(inp, prim::kPrimCellBackwardHook)) { + MS_LOG(WARNING) + << "Hook operation does not work in graph mode or ms_function, it will be eliminated during compilation."; + constexpr size_t idx = 2; + auto v_node = tuple_get_item->input(idx)->cast(); + MS_EXCEPTION_IF_NULL(v_node); + auto out_idx = GetValue(v_node->value()); + return inp->cast()->input(out_idx + 1); + } + } + return node; +} + AnfNodePtr HandleRealToComplex(const AnfNodePtr &input, const CNodePtr &din, const FuncGraphPtr &fg) { MS_EXCEPTION_IF_NULL(input); TypePtr input_type = input->Type(); @@ -244,12 +290,7 @@ void DFunctor::BackPropagate(const CNodePtr &cnode_morph, const CNodePtr &k_app, } for (size_t i = 0; i < cnode_morph->size(); i++) { auto din = tape_->NewCNode({NewValueNode(prim::kPrimTupleGetItem), bprop_app, NewValueNode(SizeToLong(i))}); - auto input = cnode_morph->input(i); - // Skip HookBackward op - if (IsPrimitiveCNode(input, prim::kPrimHookBackward)) { - auto inp_i = input->cast(); - input = inp_i->input(1); - } + auto input = SkipHookNodeInBackProp(cnode_morph->input(i)); auto din_with_real = HandleRealToComplex(input, din, tape_); MS_EXCEPTION_IF_NULL(din_with_real); din = din_with_real->cast(); @@ -299,14 +340,7 @@ AdjointPtr DFunctor::MapMorphism(const AnfNodePtr &morph) { std::vector inputs; std::vector param_adjoints; for (size_t i = 0; i < cnode_morph->size(); i++) { - auto node = cnode_morph->input(i); - // Skip HookBackward op - if (IsPrimitiveCNode(node, prim::kPrimHookBackward)) { - auto input_i = node->cast(); - MS_LOG(WARNING) - << "Hook operation does not work in graph mode or ms_function, it will be eliminated during compilation."; - node = input_i->input(1); - } + auto node = SkipHookNodeInBackProp(cnode_morph->input(i)); AdjointPtr node_adjoint = nullptr; auto node_adjoint_iter = anfnode_to_adjoin_.find(node); if (node_adjoint_iter != anfnode_to_adjoin_.end()) { @@ -457,14 +491,9 @@ void DFunctor::MapMorphism() { // Handle free morphism before output, because in some case, free morphism might depend on output's fv tangent MapFreeMorphism(); - // Skip HookBackward when it is the output node. + // Skip HookBackward op and CellBackwardHook op when it is the output node. auto output_node = primal_graph_->output(); - if (IsPrimitiveCNode(output_node, prim::kPrimHookBackward)) { - auto output_cnode = output_node->cast(); - MS_LOG(WARNING) - << "Hook operation does not work in graph mode or ms_function, it will be eliminated during compilation."; - output_node = output_cnode->input(1); - } + output_node = SkipHookNodeInBackProp(output_node); // Handle morphism from output. (void)MapMorphism(output_node); @@ -662,7 +691,9 @@ void DFunctor::MapValueObject() { if (IsValueNode(node)) { // Primitive. auto prim = GetValueNode(node); if (GetValueNode(node) == prim::kPrimReturn || - (prim->Hash() == prim::kPrimHookBackward->Hash() && prim->name() == prim::kPrimHookBackward->name())) { + (prim->Hash() == prim::kPrimHookBackward->Hash() && prim->name() == prim::kPrimHookBackward->name()) || + (prim->Hash() == prim::kPrimCellBackwardHook->Hash() && + prim->name() == prim::kPrimCellBackwardHook->name())) { continue; } MS_LOG(DEBUG) << "Map Primitive node " << node->DebugString() << "."; diff --git a/mindspore/ccsrc/frontend/optimizer/ad/dfunctor.h b/mindspore/ccsrc/frontend/optimizer/ad/dfunctor.h index b5e9b96daad..ed0ef863130 100644 --- a/mindspore/ccsrc/frontend/optimizer/ad/dfunctor.h +++ b/mindspore/ccsrc/frontend/optimizer/ad/dfunctor.h @@ -159,6 +159,8 @@ class KPrim { private: FuncGraphPtr GetBprop(const PrimitivePtr &prim, const pipeline::ResourceBasePtr &resources = nullptr); FuncGraphPtr GetFprop(const PrimitivePtr &prim); + FuncGraphPtr GetPrimBprop(const PrimitivePtr &prim, const ValueNodePtr &value_node, + const pipeline::ResourceBasePtr &resources); FuncGraphPtr FakeBprop(const ValueNodePtr &value_node, const pipeline::ResourceBasePtr &resources); FuncGraphPtr BpropCut(const ValueNodePtr &value_node, const pipeline::ResourceBasePtr &resources); // Given a bprop rule, do the K mapping. diff --git a/mindspore/ccsrc/frontend/optimizer/ad/kprim.cc b/mindspore/ccsrc/frontend/optimizer/ad/kprim.cc index bc7d11dd343..c02664cb66b 100644 --- a/mindspore/ccsrc/frontend/optimizer/ad/kprim.cc +++ b/mindspore/ccsrc/frontend/optimizer/ad/kprim.cc @@ -335,6 +335,28 @@ void KPrim::ExportBpropMindir(const py::object &obj) { } #endif +FuncGraphPtr KPrim::GetPrimBprop(const PrimitivePtr &prim, const ValueNodePtr &value_node, + const pipeline::ResourceBasePtr &resources) { + MS_EXCEPTION_IF_NULL(prim); + MS_EXCEPTION_IF_NULL(value_node); + FuncGraphPtr bprop_fg = nullptr; + auto iter = bprop_registry_.find(prim); + if (iter != bprop_registry_.end()) { + bprop_fg = iter->second; + } + + if (bprop_fg == nullptr) { + bprop_fg = GetBprop(prim, resources); + if (bprop_fg != nullptr) { + // Set bprop_g graph cache + bprop_registry_[prim] = bprop_fg; + } else { + bprop_fg = FakeBprop(value_node, resources); + } + } + return bprop_fg; +} + FuncGraphPtr KPrim::GetBprop(const PrimitivePtr &prim, const pipeline::ResourceBasePtr &resources) { // Set a child scope named "grad'PrimitiveName'" for the bprop function, // and add "Gradients" to the front. @@ -518,28 +540,16 @@ FuncGraphPtr KPrim::KPrimitive(const CNodePtr &cnode, const ValueNodePtr &value_ } FuncGraphPtr bprop_fg = nullptr; - if (prim->Hash() == prim::kPrimHookBackward->Hash() && prim->name() == prim::kPrimHookBackward->name()) { + if ((prim->Hash() == prim::kPrimHookBackward->Hash() && prim->name() == prim::kPrimHookBackward->name()) || + (prim->Hash() == prim::kPrimCellBackwardHook->Hash() && prim->name() == prim::kPrimCellBackwardHook->name())) { if (MsContext::GetInstance()->get_param(MsCtxParam::MS_CTX_EXECUTION_MODE) == kGraphMode) { MS_LOG(EXCEPTION) - << "The Primitive 'HookBackward' is not supported in graph mode, which is only supported in pynative mode.\n" + << "The Hook operation is not supported in graph mode, which is only supported in pynative mode.\n" << trace::GetDebugInfo(cnode->debug_info()); } bprop_fg = BpropCut(value_node, resources); } else { - auto iter = bprop_registry_.find(prim); - if (iter != bprop_registry_.end()) { - bprop_fg = iter->second; - } - - if (bprop_fg == nullptr) { - bprop_fg = GetBprop(prim, resources); - if (bprop_fg != nullptr) { - // Set bprop_g graph cache - bprop_registry_[prim] = bprop_fg; - } else { - bprop_fg = FakeBprop(value_node, resources); - } - } + bprop_fg = GetPrimBprop(prim, value_node, resources); } AdjustForAutoMonad(prim, bprop_fg); @@ -766,8 +776,10 @@ FuncGraphPtr KPrim::BpropCut(const ValueNodePtr &value_node, const pipeline::Res auto func_graph = std::make_shared(); std::vector outputs; + auto prim_py = prim->cast(); + MS_EXCEPTION_IF_NULL(prim_py); auto bprop_cut = std::make_shared("bprop_cut"); - bprop_cut->CopyHookFunction(prim); + bprop_cut->CopyHookFunction(prim_py); auto cell_id = GetValue(prim->GetAttr("cell_id")); if (cell_id != "") { diff --git a/mindspore/ccsrc/frontend/optimizer/ad/kpynative.cc b/mindspore/ccsrc/frontend/optimizer/ad/kpynative.cc index 7bd61fc5501..3370979c60c 100644 --- a/mindspore/ccsrc/frontend/optimizer/ad/kpynative.cc +++ b/mindspore/ccsrc/frontend/optimizer/ad/kpynative.cc @@ -398,7 +398,7 @@ bool KPynativeCellImpl::KPynativeOp(const CNodePtr &cnode, const ValuePtrList &o } FuncGraphPtr bprop_fg = nullptr; - if (IsPrimitiveEquals(prim, prim::kPrimHookBackward)) { + if (IsPrimitiveEquals(prim, prim::kPrimHookBackward) || IsPrimitiveEquals(prim, prim::kPrimCellBackwardHook)) { bprop_fg = BuildBPropCutFuncGraph(prim, cnode); } else if (IsPrimitiveEquals(prim, prim::kPrimMakeTuple) || IsPrimitiveEquals(prim, prim::kPrimMakeList)) { bprop_fg = BuildMakeSequenceBprop(prim, cnode); @@ -906,8 +906,11 @@ FuncGraphPtr KPynativeCellImpl::BuildBPropCutFuncGraph(const PrimitivePtr &prim, auto func_graph = std::make_shared(); std::vector outputs; + auto prim_py = prim->cast(); + MS_EXCEPTION_IF_NULL(prim_py); auto bprop_cut = std::make_shared("bprop_cut"); - bprop_cut->CopyHookFunction(prim); + bprop_cut->CopyHookFunction(prim_py); + prim_py->AddBpropCutPrim(bprop_cut); auto cell_id = GetValue(prim->GetAttr("cell_id")); if (cell_id != "") { diff --git a/mindspore/ccsrc/frontend/optimizer/ad/prim_bprop_optimizer.cc b/mindspore/ccsrc/frontend/optimizer/ad/prim_bprop_optimizer.cc index 847339f2adb..27027e6856a 100644 --- a/mindspore/ccsrc/frontend/optimizer/ad/prim_bprop_optimizer.cc +++ b/mindspore/ccsrc/frontend/optimizer/ad/prim_bprop_optimizer.cc @@ -194,7 +194,8 @@ FuncGraphPtr PrimBpropOptimizer::OptimizeBPropFuncGraph(const FuncGraphPtr &bpro MS_LOG(DEBUG) << "Hash of prim " << prim->ToString() << " is:" << prim->hash(); // kPrimHookBackward - bool hookback_flg = IsPrimitiveEquals(prim, prim::kPrimHookBackward); + bool hookback_flg = + IsPrimitiveEquals(prim, prim::kPrimHookBackward) || IsPrimitiveEquals(prim, prim::kPrimCellBackwardHook); if (hookback_flg || IsPrimitiveEquals(prim, prim::kPrimMakeTuple) || IsPrimitiveEquals(prim, prim::kPrimMakeList)) { return GenSpecOptBprop(bprop_fg, op_args, out, prim, hookback_flg); } diff --git a/mindspore/ccsrc/frontend/optimizer/irpass.cc b/mindspore/ccsrc/frontend/optimizer/irpass.cc index f8f959636de..1d47bf958d6 100644 --- a/mindspore/ccsrc/frontend/optimizer/irpass.cc +++ b/mindspore/ccsrc/frontend/optimizer/irpass.cc @@ -61,10 +61,10 @@ OptimizeIRPassLib::OptimizeIRPassLib() { prim::kPrimIdentity, prim::kPrimMomentum, prim::kPrimMul, prim::kPrimPow}); arithmetic_simplify2_ = MakeSubstitution(std::make_shared(), "arithmetic_simplify2", {prim::kPrimMul}); - special_op_eliminate_ = - MakeSubstitution(std::make_shared(), "special_op_eliminate", - {prim::kPrimInsertGradientOf, prim::kPrimStopGradient, prim::kPrimHookBackward, - prim::kPrimPrintShapeType, prim::kPrimGetRefValue, prim::kPrimMirror, prim::kPrimVirtualDiv}); + special_op_eliminate_ = MakeSubstitution( + std::make_shared(), "special_op_eliminate", + {prim::kPrimInsertGradientOf, prim::kPrimStopGradient, prim::kPrimHookBackward, prim::kPrimCellBackwardHook, + prim::kPrimPrintShapeType, prim::kPrimGetRefValue, prim::kPrimMirror, prim::kPrimVirtualDiv}); pynative_eliminate_ = MakeSubstitution(std::make_shared(), "pynative_eliminate", IsCNodeDup); zero_like_fill_zero_ = MakeSubstitution(std::make_shared(), "zero_like_fill_zero", prim::kPrimZerosLike); diff --git a/mindspore/ccsrc/frontend/optimizer/irpass/special_op_eliminate.h b/mindspore/ccsrc/frontend/optimizer/irpass/special_op_eliminate.h index 2c2f26a4086..066e5508cb9 100644 --- a/mindspore/ccsrc/frontend/optimizer/irpass/special_op_eliminate.h +++ b/mindspore/ccsrc/frontend/optimizer/irpass/special_op_eliminate.h @@ -45,6 +45,7 @@ class SpecialOpEliminater : public OptimizerCaller { : insert_gradient_of_(std::make_shared(prim::kPrimInsertGradientOf)), stop_gradient_(std::make_shared(prim::kPrimStopGradient)), hook_backward_(std::make_shared(prim::kPrimHookBackward)), + cell_backward_hook_(std::make_shared(prim::kPrimCellBackwardHook)), print_shape_type_(std::make_shared(prim::kPrimPrintShapeType)), get_ref_value_(std::make_shared(prim::kPrimGetRefValue)), mirror_(std::make_shared(prim::kPrimMirror)), @@ -52,6 +53,7 @@ class SpecialOpEliminater : public OptimizerCaller { eliminaters_.emplace_back(insert_gradient_of_); eliminaters_.emplace_back(stop_gradient_); eliminaters_.emplace_back(hook_backward_); + eliminaters_.emplace_back(cell_backward_hook_); eliminaters_.emplace_back(print_shape_type_); eliminaters_.emplace_back(get_ref_value_); eliminaters_.emplace_back(mirror_); @@ -64,7 +66,7 @@ class SpecialOpEliminater : public OptimizerCaller { for (auto &eliminater : eliminaters_) { new_node = (*eliminater)(optimizer, node); if (new_node != nullptr) { - if (IsPrimitiveCNode(node, prim::kPrimHookBackward)) { + if (IsPrimitiveCNode(node, prim::kPrimHookBackward) || IsPrimitiveCNode(node, prim::kPrimCellBackwardHook)) { MS_LOG(WARNING) << "Hook operation does not work in graph mode or ms_function, it will be eliminated during compilation."; } @@ -75,8 +77,8 @@ class SpecialOpEliminater : public OptimizerCaller { } private: - OptimizerCallerPtr insert_gradient_of_, stop_gradient_, hook_backward_, print_shape_type_, get_ref_value_, mirror_, - virtual_div_; + OptimizerCallerPtr insert_gradient_of_, stop_gradient_, hook_backward_, cell_backward_hook_, print_shape_type_, + get_ref_value_, mirror_, virtual_div_; std::vector eliminaters_{}; }; diff --git a/mindspore/ccsrc/pipeline/jit/parse/data_converter.cc b/mindspore/ccsrc/pipeline/jit/parse/data_converter.cc index 90719e7a958..ec2e00fb964 100644 --- a/mindspore/ccsrc/pipeline/jit/parse/data_converter.cc +++ b/mindspore/ccsrc/pipeline/jit/parse/data_converter.cc @@ -141,7 +141,7 @@ FuncGraphPtr ConvertToBpropCut(const py::object &obj) { std::vector outputs; auto fake_bprop = std::make_shared("bprop_cut"); - fake_bprop->set_hook(bprop_func); + fake_bprop->AddBackwardHookFn(0, bprop_func); (void)fake_bprop->AddAttr(CUSTOM_BPROP_NAME, MakeValue(true)); outputs.push_back(NewValueNode(fake_bprop)); diff --git a/mindspore/ccsrc/pipeline/pynative/pynative_execute.cc b/mindspore/ccsrc/pipeline/pynative/pynative_execute.cc index 77185e3219b..d4486bccc08 100644 --- a/mindspore/ccsrc/pipeline/pynative/pynative_execute.cc +++ b/mindspore/ccsrc/pipeline/pynative/pynative_execute.cc @@ -75,8 +75,8 @@ const size_t ARG_SIZE = 2; const size_t MAX_TOP_CELL_COUNTS = 20; // primitive unable to infer value for constant input in PyNative mode -const std::set kVmOperators = {"make_ref", "HookBackward", "InsertGradientOf", "stop_gradient", - "mixed_precision_cast"}; +const std::set kVmOperators = {"make_ref", "InsertGradientOf", "stop_gradient", "mixed_precision_cast", + "HookBackward", "CellBackwardHook"}; const char kOpsFunctionModelName[] = "mindspore.ops.functional"; const char kGrad[] = "grad"; std::map> kSessionBackends; @@ -1728,6 +1728,22 @@ void GradExecutor::EnableOpGraphCache(bool is_enable) { inst->set_param(MS_CTX_ENABLE_PYNATIVE_OP_GRAPH_CACHE, is_enable); } +void GradExecutor::SetHookChanged(const py::object &cell) { + auto cell_id = GetId(cell); + for (const auto &top_cell : top_cell_list_) { + MS_EXCEPTION_IF_NULL(top_cell); + if (top_cell->cell_id().find(cell_id) != std::string::npos) { + top_cell->set_hook_changed(true); + } + const auto &sub_cells = top_cell->sub_cell_list(); + for (const auto &sub_cell_id : sub_cells) { + if (sub_cell_id.find(cell_id) != std::string::npos) { + top_cell->set_hook_changed(true); + } + } + } +} + void GradExecutor::RecordGradOpInfo(const OpExecInfoPtr &op_exec_info) { if (!grad_flag_) { MS_LOG(DEBUG) << "Grad flag is set to false, no need to record op info"; @@ -2050,14 +2066,17 @@ py::object ForwardExecutor::RunOpInVM(const OpExecInfoPtr &op_exec_info) { MS_EXCEPTION_IF_NULL(op_exec_info->py_primitive); auto &op_inputs = op_exec_info->op_inputs; - if (op_exec_info->op_name == "HookBackward" || op_exec_info->op_name == "InsertGradientOf" || - op_exec_info->op_name == "stop_gradient") { + if (op_exec_info->op_name == prim::kPrimInsertGradientOf->name() || + op_exec_info->op_name == prim::kPrimStopGradient->name() || + op_exec_info->op_name == prim::kPrimHookBackward->name() || + op_exec_info->op_name == prim::kPrimCellBackwardHook->name()) { py::tuple result(op_inputs.size()); for (size_t i = 0; i < op_inputs.size(); i++) { py::object input = op_inputs[i]; auto tensor = py::cast(input); MS_EXCEPTION_IF_NULL(tensor); - if (op_exec_info->op_name == "HookBackward") { + if (op_exec_info->op_name == prim::kPrimHookBackward->name() || + op_exec_info->op_name == prim::kPrimCellBackwardHook->name()) { // the input object is not a output of forward cnode, eg: parameter result[i] = tensor; } else { @@ -2405,7 +2424,10 @@ void GradExecutor::NewGraphInner(py::object *ret, const py::object &cell, const // Top cell forward run. const auto &pre_top_cell = top_it->second; MS_EXCEPTION_IF_NULL(pre_top_cell); - if (!pre_top_cell->is_dynamic()) { + if (pre_top_cell->hook_changed()) { + already_run_top_cell_.erase(top_it); + EraseTopCellFromTopCellList(pre_top_cell); + } else if (!pre_top_cell->is_dynamic()) { MS_LOG(DEBUG) << "Top cell " << cell_id << " is not dynamic, no need to run NewGraphInner again"; ResetTopCellInfo(pre_top_cell, args); PushHighOrderGraphStack(pre_top_cell); @@ -2594,7 +2616,7 @@ void GradExecutor::DoGradForCustomBprop(const py::object &cell, const py::object auto cell_ptr = py::cast(cell); fake_prim->set_bprop_cls_name(cell_ptr->name()); } - fake_prim->set_hook(bprop_func); + fake_prim->AddBackwardHookFn(0, bprop_func); const auto &cell_id = GetCellId(cell, args); (void)fake_prim->AddAttr("cell_id", MakeValue(cell_id)); @@ -2911,7 +2933,7 @@ py::object GradExecutor::CheckAlreadyRun(const prim::GradOperationPtr &grad, con auto find_top_cell = GetTopCell(check_already_run_cell_id); if (find_top_cell != nullptr) { MS_LOG(DEBUG) << "Find already run top cell"; - forward_run = find_top_cell->forward_already_run(); + forward_run = find_top_cell->forward_already_run() && !find_top_cell->hook_changed(); auto curr_top_cell = top_cell(); set_top_cell(find_top_cell); bool input_args_changed = @@ -3293,6 +3315,13 @@ bool PynativeExecutor::grad_flag() const { return grad_executor()->grad_flag(); void PynativeExecutor::set_grad_flag(bool flag) { grad_executor()->set_grad_flag(flag); } +void PynativeExecutor::SetHookChanged(const py::object &cell) { + if (!py::isinstance(cell)) { + MS_LOG(EXCEPTION) << "The 'set_hook_changed' function is only supported on Cell object!"; + } + grad_executor()->SetHookChanged(cell); +} + void PynativeExecutor::set_graph_phase(const std::string &graph_phase) { grad_executor()->set_graph_phase(graph_phase); } @@ -3475,6 +3504,7 @@ REGISTER_PYBIND_DEFINE(PynativeExecutor_, ([](const py::module *m) { .def("__call__", &PynativeExecutor::Run, "pynative executor run grad graph.") .def("set_graph_phase", &PynativeExecutor::set_graph_phase, "pynative set graph phase") .def("grad_flag", &PynativeExecutor::grad_flag, "pynative grad flag") + .def("set_hook_changed", &PynativeExecutor::SetHookChanged, "set pynative hook changed") .def("set_grad_position", &PynativeExecutor::set_grad_position, "set pynative grad position") .def("set_grad_flag", &PynativeExecutor::set_grad_flag, py::arg("flag") = py::bool_(false), "Executor set grad flag.") diff --git a/mindspore/ccsrc/pipeline/pynative/pynative_execute.h b/mindspore/ccsrc/pipeline/pynative/pynative_execute.h index 34f382d4b1a..bc0e5a3a2cf 100644 --- a/mindspore/ccsrc/pipeline/pynative/pynative_execute.h +++ b/mindspore/ccsrc/pipeline/pynative/pynative_execute.h @@ -78,6 +78,8 @@ class TopCellInfo { void set_grad_order(size_t grad_order) { grad_order_ = grad_order; } bool is_dynamic() const { return is_dynamic_; } void set_is_dynamic(bool is_dynamic) { is_dynamic_ = is_dynamic; } + bool hook_changed() const { return hook_changed_; } + void set_hook_changed(bool hook_changed) { hook_changed_ = hook_changed; } bool vm_compiled() const { return vm_compiled_; } void set_vm_compiled(bool vm_compiled) { vm_compiled_ = vm_compiled; } bool ms_function_flag() const { return ms_function_flag_; } @@ -123,6 +125,7 @@ class TopCellInfo { bool is_topest_{false}; bool is_dynamic_{false}; bool vm_compiled_{false}; + bool hook_changed_{false}; bool ms_function_flag_{false}; bool is_init_kpynative_{false}; bool forward_already_run_{false}; @@ -191,6 +194,7 @@ class GradExecutor { size_t GetHighOrderStackSize() const { return high_order_stack_.size(); } TopCellInfoPtr GetTopCell(const string &already_run_cell_id); void EnableOpGraphCache(bool is_enable); + void SetHookChanged(const py::object &cell); bool need_renormalize() const { return need_renormalize_; } bool enable_op_cache() const { return enable_op_cache_; } bool grad_is_running() const { return grad_is_running_; } @@ -397,6 +401,7 @@ class PynativeExecutor : public std::enable_shared_from_this { void set_graph_phase(const std::string &graph_phase); void set_py_exe_path(const py::object &py_exe_path); void set_kernel_build_server_dir(const py::object &kernel_build_server_dir); + void SetHookChanged(const py::object &cell); void NewGraph(const py::object &cell, const py::args &args); void EndGraph(const py::object &cell, const py::object &out, const py::args &args); void GradNet(const prim::GradOperationPtr &grad, const py::object &cell, const py::object &weights, diff --git a/mindspore/ccsrc/pybind_api/ir/primitive_py.cc b/mindspore/ccsrc/pybind_api/ir/primitive_py.cc index 144b3bca390..8ad9f9cf0b5 100644 --- a/mindspore/ccsrc/pybind_api/ir/primitive_py.cc +++ b/mindspore/ccsrc/pybind_api/ir/primitive_py.cc @@ -55,6 +55,54 @@ void SyncData(const py::object &arg) { tensor->data_sync(); } } + +void ConvertCTensorToPyTensor(const py::tuple &input_args, py::tuple *convert_args) { + MS_EXCEPTION_IF_NULL(convert_args); + if (input_args.size() != (*convert_args).size()) { + MS_LOG(EXCEPTION) << "The size of input_args: " << input_args.size() + << " should be equal to the size of convert_args: " << (*convert_args).size(); + } + for (size_t i = 0; i < input_args.size(); ++i) { + if (py::isinstance(input_args[i])) { + (*convert_args)[i] = parse::python_adapter::CallPyFn(parse::PYTHON_MOD_PARSE_MODULE, + parse::PYTHON_MOD_CONVERT_TO_MS_TENSOR, input_args[i]); + } else if (py::isinstance(input_args[i])) { + auto tuple_inp_arg = py::cast(input_args[i]); + py::tuple convert_tuple_arg(tuple_inp_arg.size()); + ConvertCTensorToPyTensor(tuple_inp_arg, &convert_tuple_arg); + (*convert_args)[i] = convert_tuple_arg; + } else { + (*convert_args)[i] = input_args[i]; + } + } +} + +py::tuple ConstructCellHookFnArgs(const std::string &cell_id, const py::object &grad_input, + const py::object &grad_output) { + constexpr size_t grad_input_index = 1; + constexpr size_t grad_output_index = 2; + constexpr size_t input_args_nums = 3; + // Convert c++ object to python object. + py::tuple c_grad_args(input_args_nums - 1); + c_grad_args[0] = grad_input; + c_grad_args[1] = grad_output; + py::tuple py_grad_args(input_args_nums - 1); + ConvertCTensorToPyTensor(c_grad_args, &py_grad_args); + // Get tuple args of cell hook function. + py::tuple hook_fn_args(input_args_nums); + hook_fn_args[0] = cell_id; + if (!py::isinstance(py_grad_args[0])) { + hook_fn_args[grad_input_index] = py::make_tuple(py_grad_args[0]); + } else { + hook_fn_args[grad_input_index] = py_grad_args[0]; + } + if (!py::isinstance(py_grad_args[1])) { + hook_fn_args[grad_output_index] = py::make_tuple(py_grad_args[1]); + } else { + hook_fn_args[grad_output_index] = py_grad_args[1]; + } + return hook_fn_args; +} } // namespace std::map PrimitivePy::hook_grad_; @@ -68,7 +116,9 @@ PrimitivePy::PrimitivePy(const py::object &python_obj, const PrimitivePyAdapterP Primitive::set_prim_type(adapter->prim_type_); Primitive::set_const_prim(adapter->is_const_prim_); Primitive::set_const_input_indexes(adapter->const_input_indexes_); - set_hook(adapter->hook_); + for (const auto &elem : adapter->backward_hook_fn_) { + AddBackwardHookFn(elem.first, elem.second); + } set_instance_name(adapter->instance_name_); } PrimitivePy::~PrimitivePy() {} @@ -91,7 +141,9 @@ py::function PrimitivePy::GetBpropFunction() { py::tuple check_bprop_out(const py::object &grads_obj, const py::tuple &py_args, const std::string &bprop_cls_name) { py::tuple grads; - if (!py::isinstance(grads_obj)) { + if (py::isinstance(grads_obj)) { + MS_EXCEPTION(TypeError) << "The 'grads_obj' is none."; + } else if (!py::isinstance(grads_obj)) { grads = py::make_tuple(grads_obj); } else { grads = py::cast(grads_obj); @@ -138,21 +190,37 @@ py::tuple check_bprop_out(const py::object &grads_obj, const py::tuple &py_args, return grads; } -void PrimitivePy::ConvertCTensorToPyTensor(const py::tuple &input_args, py::tuple *convert_args) const { - MS_EXCEPTION_IF_NULL(convert_args); - if (input_args.size() != (*convert_args).size()) { - MS_LOG(EXCEPTION) << "The size of input_args: " << input_args.size() - << " should be equal to the size of convert_args: " << (*convert_args).size(); - } - for (size_t i = 0; i < input_args.size(); ++i) { - (*convert_args)[i] = py::isinstance(input_args[i]) - ? parse::python_adapter::CallPyFn(parse::PYTHON_MOD_PARSE_MODULE, - parse::PYTHON_MOD_CONVERT_TO_MS_TENSOR, input_args[i]) - : input_args[i]; +void PrimitivePy::AddBpropCutPrim(const PrimitivePyPtr &bprop_cut_prim) { + MS_EXCEPTION_IF_NULL(bprop_cut_prim); + bprop_cut_prims_.emplace_back(bprop_cut_prim); +} + +void PrimitivePy::AddBackwardHookFn(const int &key, const py::function &backward_hook_fn) { + backward_hook_fn_[key] = backward_hook_fn; + for (const auto &elem : bprop_cut_prims_) { + PrimitivePyPtr bprop_cut_prim = elem.lock(); + if (bprop_cut_prim != nullptr) { + bprop_cut_prim->AddBackwardHookFn(key, backward_hook_fn); + } } } -void PrimitivePy::CheckHookConsistency(const py::object &grad_out, const py::object &expected_grad_out) const { +void PrimitivePy::RemoveBackwardHookFn(const int &key) { + auto iter = backward_hook_fn_.find(key); + if (iter != backward_hook_fn_.end()) { + backward_hook_fn_.erase(key); + } + // Remove hook_fn for bprop cut prim on grad graph. + for (const auto &elem : bprop_cut_prims_) { + PrimitivePyPtr bprop_cut_prim = elem.lock(); + if (bprop_cut_prim != nullptr) { + bprop_cut_prim->RemoveBackwardHookFn(key); + } + } +} + +void PrimitivePy::CheckHookConsistency(const py::object &grad_out, const py::object &expected_grad_out, + const py::object &code_obj, const py::object &co_name) const { if (py::isinstance(expected_grad_out)) { if (!py::isinstance(grad_out)) { hook_grad_.clear(); @@ -166,15 +234,13 @@ void PrimitivePy::CheckHookConsistency(const py::object &grad_out, const py::obj << ", but it is " << actual_out_tuple.size(); } for (size_t i = 0; i < expected_out_tuple.size(); ++i) { - CheckHookConsistency(actual_out_tuple[i], expected_out_tuple[i]); + CheckHookConsistency(actual_out_tuple[i], expected_out_tuple[i], code_obj, co_name); } } if (py::isinstance(expected_grad_out)) { if (!py::isinstance(grad_out)) { hook_grad_.clear(); - py::object code_obj = py::getattr(hook_, "__code__"); - py::object co_name = py::getattr(code_obj, "co_name"); MS_EXCEPTION(TypeError) << "The output type of:" << py::str(co_name) << " should be a tensor but got " << py::cast(grad_out.attr("__class__").attr("__name__")) << "."; } @@ -184,8 +250,6 @@ void PrimitivePy::CheckHookConsistency(const py::object &grad_out, const py::obj MS_EXCEPTION_IF_NULL(expected_out_tensor); if (actual_out_tensor->GetShapeAndDataTypeInfo() != expected_out_tensor->GetShapeAndDataTypeInfo()) { hook_grad_.clear(); - py::object code_obj = py::getattr(hook_, "__code__"); - py::object co_name = py::getattr(code_obj, "co_name"); MS_EXCEPTION(ValueError) << "The output type of " << py::str(co_name) << " is not consistent with the expected, it should be " << expected_out_tensor->GetShapeAndDataTypeInfo() << ", but got " @@ -195,24 +259,31 @@ void PrimitivePy::CheckHookConsistency(const py::object &grad_out, const py::obj } BaseRef PrimitivePy::RunCellBpropFunction(const py::tuple &py_args) const { + if (backward_hook_fn_.size() > 1) { + MS_LOG(EXCEPTION) << "Multiple registration of bprop function is not supported."; + } SyncData(py_args); - auto size = py_args.size(); - constexpr size_t grad_param_nums = 2; - py::tuple input_args(size - grad_param_nums); - for (size_t i = 0; i < size - grad_param_nums; ++i) { + py::tuple converted_args(py_args.size()); + ConvertCTensorToPyTensor(py_args, &converted_args); + constexpr size_t non_inp_args_size = 2; // out and dout. + auto inp_args_size = py_args.size() - non_inp_args_size; + py::tuple input_args(inp_args_size); + for (size_t i = 0; i < inp_args_size; ++i) { input_args[i] = py_args[i]; } - py::tuple convert_args(py_args.size()); - ConvertCTensorToPyTensor(py_args, &convert_args); + // Run bprop function. auto inst = pynative::PynativeExecutor::GetInstance(); MS_EXCEPTION_IF_NULL(inst); try { - MS_LOG(DEBUG) << "Run bprop function start"; - inst->NewGraph(hook_, input_args.cast()); - py::object grads_obj = hook_(*convert_args); - py::tuple grads = check_bprop_out(grads_obj, py_args, bprop_cls_name_); - inst->EndGraph(hook_, grads_obj, input_args.cast()); - MS_LOG(DEBUG) << "Run bprop function end"; + MS_LOG(DEBUG) << "Run bprop function start."; + py::tuple grads; + for (const auto &elem : backward_hook_fn_) { + inst->NewGraph(elem.second, input_args.cast()); + py::object grads_obj = elem.second(*converted_args); + grads = check_bprop_out(grads_obj, py_args, bprop_cls_name_); + inst->EndGraph(elem.second, grads_obj, input_args.cast()); + } + MS_LOG(DEBUG) << "Run bprop function end."; return std::make_shared(grads); } catch (std::exception &bt) { inst->ClearRes(); @@ -221,62 +292,61 @@ BaseRef PrimitivePy::RunCellBpropFunction(const py::tuple &py_args) const { } BaseRef PrimitivePy::RunCellHookFunction(const py::tuple &py_args) const { - constexpr size_t grad_input_index = 1; - constexpr size_t grad_output_index = 2; - constexpr size_t input_param_nums = 3; - SyncData(py_args[grad_output_index]); - - py::object obj; + // Get the gradient passed to current bprop cut op. + const auto args_size = py_args.size(); + py::object grad_output = py_args[args_size - 1]; + // Get the cell id. auto cell_id = GetValue(this->GetAttr(kCellIDAttrName)); auto iter = hook_grad_.find(cell_id); if (iter != hook_grad_.end()) { - py::object code_obj = py::getattr(hook_, "__code__"); - py::object co_name = py::getattr(code_obj, "co_name"); - if (std::string(py::str(co_name)) == "staging_specialize") { - py::object name_obj = py::getattr(hook_, "__name__"); - MS_LOG(EXCEPTION) << "Decorating hook function " << py::str(name_obj) << " with '@ms_function' is not supported."; + // The second bprop_cut used to hook output gradient of cell. + for (const auto &elem : backward_hook_fn_) { + py::object code_obj = py::getattr(elem.second, "__code__"); + py::object co_name = py::getattr(code_obj, "co_name"); + if (std::string(py::str(co_name)) == "staging_specialize") { + py::object name_obj = py::getattr(elem.second, "__name__"); + MS_LOG(EXCEPTION) << "Decorating hook function " << py::str(name_obj) + << " with '@ms_function' is not supported."; + } + SyncData(grad_output); + py::tuple hook_fn_args = ConstructCellHookFnArgs(cell_id, iter->second, grad_output); + py::object ret = elem.second(*hook_fn_args); + if (!py::isinstance(ret)) { + grad_output = ret; + } + CheckHookConsistency(grad_output, py_args[args_size - 1], code_obj, co_name); } - - py::tuple convert_args(input_param_nums - 1); - py::tuple input_args(input_param_nums - 1); - input_args[0] = iter->second; - input_args[1] = py_args[grad_output_index]; - ConvertCTensorToPyTensor(input_args, &convert_args); - auto hook_args = py::tuple(input_param_nums); - hook_args[0] = cell_id; - hook_args[grad_input_index] = py::make_tuple(convert_args[0]); - hook_args[grad_output_index] = py::make_tuple(convert_args[1]); - obj = hook_(*hook_args); - if (py::isinstance(obj)) { - obj = py_args[grad_output_index]; - } - CheckHookConsistency(obj, py_args[grad_output_index]); (void)hook_grad_.erase(cell_id); } else { - hook_grad_[cell_id] = py_args[grad_output_index]; - obj = py_args[grad_output_index]; + // The first bprop_cut used to hook input gradient of cell. + SyncData(grad_output); + hook_grad_[cell_id] = grad_output; } - obj = py::make_tuple(obj); - return std::make_shared(obj); + if (!py::isinstance(grad_output)) { + grad_output = py::make_tuple(grad_output); + } + return std::make_shared(grad_output); } BaseRef PrimitivePy::RunVariableHookFunction(const py::tuple &py_args) const { - py::object code_obj = py::getattr(hook_, "__code__"); - py::object co_name = py::getattr(code_obj, "co_name"); - if (std::string(py::str(co_name)) == "staging_specialize") { - py::object name_obj = py::getattr(hook_, "__name__"); - MS_LOG(EXCEPTION) << "Decorating hook function " << py::str(name_obj) << " with '@ms_function' is not supported."; - } - constexpr size_t grad_output_index = 2; - SyncData(py_args[grad_output_index]); - py::object obj = hook_(py::make_tuple(py_args[grad_output_index])); - if (py::isinstance(obj)) { - obj = py_args[grad_output_index]; + py::object grad_output = py_args[grad_output_index]; + for (const auto &elem : backward_hook_fn_) { + py::object code_obj = py::getattr(elem.second, "__code__"); + py::object co_name = py::getattr(code_obj, "co_name"); + if (std::string(py::str(co_name)) == "staging_specialize") { + py::object name_obj = py::getattr(elem.second, "__name__"); + MS_LOG(EXCEPTION) << "Decorating hook function " << py::str(name_obj) << " with '@ms_function' is not supported."; + } + SyncData(grad_output); + py::object ret = elem.second(py::make_tuple(grad_output)); + if (!py::isinstance(ret)) { + grad_output = ret; + } + CheckHookConsistency(grad_output, py_args[grad_output_index], code_obj, co_name); } - CheckHookConsistency(obj, py_args[grad_output_index]); - obj = py::make_tuple(obj); - return std::make_shared(obj); + grad_output = py::make_tuple(grad_output); + return std::make_shared(grad_output); } BaseRef PrimitivePy::RunHookFunction(const VectorRef &args) const { @@ -321,14 +391,12 @@ py::dict PrimitivePy::GetAttrDict() { return attr_dict; } -void PrimitivePy::CopyHookFunction(const PrimitivePtr &primitive) { - MS_EXCEPTION_IF_NULL(primitive); - if (!primitive->isa()) { - MS_LOG(EXCEPTION) << "Cannot copy a primitive which is not python primitive hook function to python primitive!"; - } - auto primitive_py = primitive->cast(); +void PrimitivePy::CopyHookFunction(const PrimitivePyPtr &primitive_py) { MS_EXCEPTION_IF_NULL(primitive_py); - this->set_hook(primitive_py->hook()); + const auto &backward_hook_fn = primitive_py->backward_hook_fn(); + for (const auto &elem : backward_hook_fn) { + AddBackwardHookFn(elem.first, elem.second); + } if (primitive_py->HasAttr(kBpropAttrName)) { set_bprop_cls_name(primitive_py->bprop_cls_name_); (void)this->AddAttr(kBpropAttrName, primitive_py->GetAttr(kBpropAttrName)); @@ -500,11 +568,24 @@ void PrimitivePyAdapter::set_signatures(const std::vector &signatures } } -void PrimitivePyAdapter::set_hook(const py::function &hook) { - hook_ = hook; +int PrimitivePyAdapter::AddBackwardHookFn(const py::function &backward_hook_fn) { + ++backward_hook_fn_key_; + backward_hook_fn_[backward_hook_fn_key_] = backward_hook_fn; auto prim = attached_primitive_.lock(); if (prim != nullptr) { - prim->set_hook(hook); + prim->AddBackwardHookFn(backward_hook_fn_key_, backward_hook_fn); + } + return backward_hook_fn_key_; +} + +void PrimitivePyAdapter::RemoveBackwardHookFn(int key) { + auto iter = backward_hook_fn_.find(key); + if (iter != backward_hook_fn_.end()) { + backward_hook_fn_.erase(iter); + } + auto prim = attached_primitive_.lock(); + if (prim != nullptr) { + prim->RemoveBackwardHookFn(key); } } @@ -524,27 +605,28 @@ void PrimitivePyAdapter::set_attached_primitive(const PrimitivePyPtr &prim) { attached_primitive_ = prim; } -REGISTER_PYBIND_DEFINE(Primitive_, ([](const py::module *m) { - (void)py::enum_(*m, "prim_type", py::arithmetic()) - .value("unknown", PrimType::kPrimTypeUnknown) - .value("builtin", PrimType::kPrimTypeBuiltIn) - .value("py_infer_shape", PrimType::kPrimTypePyInfer) - .value("user_custom", PrimType::kPrimTypeUserCustom) - .value("py_infer_check", PrimType::kPrimTypePyCheck); - (void)py::class_>(*m, "Primitive_") - .def_readonly(PYTHON_PRIMITIVE_FLAG, &PrimitivePyAdapter::parse_info_) - .def(py::init()) - .def("add_attr", &PrimitivePyAdapter::AddPyAttr, "add primitive attr") - .def("del_attr", &PrimitivePyAdapter::DelPyAttr, "del primitive attr") - .def("get_attr_dict", &PrimitivePyAdapter::GetAttrDict, "get primitive attr") - .def("set_prim_type", &PrimitivePyAdapter::set_prim_type, "Set primitive type.") - .def("set_const_prim", &PrimitivePyAdapter::set_const_prim, "Set primitive is const.") - .def("set_const_input_indexes", &PrimitivePyAdapter::set_const_input_indexes, - "Set primitive const input indexes.") - .def("set_signatures", &PrimitivePyAdapter::set_signatures, - "Set primitive inputs signature.") - .def("register_hook", &PrimitivePyAdapter::set_hook, "Set primitive hook function.") - .def("set_instance_name", &PrimitivePyAdapter::set_instance_name, - "Set primitive instance name."); - })); +REGISTER_PYBIND_DEFINE( + Primitive_, ([](const py::module *m) { + (void)py::enum_(*m, "prim_type", py::arithmetic()) + .value("unknown", PrimType::kPrimTypeUnknown) + .value("builtin", PrimType::kPrimTypeBuiltIn) + .value("py_infer_shape", PrimType::kPrimTypePyInfer) + .value("user_custom", PrimType::kPrimTypeUserCustom) + .value("py_infer_check", PrimType::kPrimTypePyCheck); + (void)py::class_>(*m, "Primitive_") + .def_readonly(PYTHON_PRIMITIVE_FLAG, &PrimitivePyAdapter::parse_info_) + .def(py::init()) + .def("add_attr", &PrimitivePyAdapter::AddPyAttr, "add primitive attr") + .def("del_attr", &PrimitivePyAdapter::DelPyAttr, "del primitive attr") + .def("get_attr_dict", &PrimitivePyAdapter::GetAttrDict, "get primitive attr") + .def("set_prim_type", &PrimitivePyAdapter::set_prim_type, "Set primitive type.") + .def("set_const_prim", &PrimitivePyAdapter::set_const_prim, "Set primitive is const.") + .def("set_const_input_indexes", &PrimitivePyAdapter::set_const_input_indexes, + "Set primitive const input indexes.") + .def("set_signatures", &PrimitivePyAdapter::set_signatures, "Set primitive inputs signature.") + .def("add_backward_hook_fn", &PrimitivePyAdapter::AddBackwardHookFn, "Add primitive backward hook function.") + .def("remove_backward_hook_fn", &PrimitivePyAdapter::RemoveBackwardHookFn, + "Remove primitive backward hook function.") + .def("set_instance_name", &PrimitivePyAdapter::set_instance_name, "Set primitive instance name."); + })); } // namespace mindspore diff --git a/mindspore/ccsrc/pybind_api/ir/primitive_py.h b/mindspore/ccsrc/pybind_api/ir/primitive_py.h index 5701cea41b2..cfd00910896 100644 --- a/mindspore/ccsrc/pybind_api/ir/primitive_py.h +++ b/mindspore/ccsrc/pybind_api/ir/primitive_py.h @@ -48,17 +48,15 @@ class PrimitivePy : public Primitive { PrimitivePy(const py::object &python_obj, const PrimitivePyAdapterPtr &adapter); ~PrimitivePy() override; MS_DECLARE_PARENT(PrimitivePy, Primitive); + const bool parse_info_ = true; py::function GetBpropFunction(); - void set_signatures(const std::vector &signatures); - const std::vector &signatures() const { return signatures_; } - - void CopyHookFunction(const PrimitivePtr &primitive); - - py::dict GetAttrDict(); - void set_hook(const py::function &hook) { hook_ = hook; } - py::function hook() const { return hook_; } + const std::map &backward_hook_fn() const { return backward_hook_fn_; } + void CopyHookFunction(const PrimitivePyPtr &primitive_py); + void AddBpropCutPrim(const PrimitivePyPtr &bprop_cut_prim); + void AddBackwardHookFn(const int &key, const py::function &backward_hook_fn); + void RemoveBackwardHookFn(const int &key); BaseRef RunHookFunction(const VectorRef &args) const; BaseRef RunCellBpropFunction(const py::tuple &py_args) const; BaseRef RunCellHookFunction(const py::tuple &py_args) const; @@ -66,12 +64,12 @@ class PrimitivePy : public Primitive { BaseRef RunComputeFunction(const VectorRef &args) const override; py::object RunPyComputeFunction(const py::tuple &py_args) const; bool HasComputeFunction() const; - const bool parse_info_ = true; + py::dict GetAttrDict(); const py::object &GetPyObj() const { return python_obj_; } - py::dict RunInfer(const py::tuple &args); - void RunCheck(const py::tuple &args); - py::object RunInferValue(const py::tuple &args); bool HasPyObj() const { return python_obj_.operator bool(); } + void RunCheck(const py::tuple &args); + py::dict RunInfer(const py::tuple &args); + py::object RunInferValue(const py::tuple &args); PrimitivePtr Clone() override; PrimitivePyAdapterPtr adapter() const { return adapter_; } void set_bprop_cls_name(const std::string &name) { bprop_cls_name_ = name; } @@ -79,13 +77,14 @@ class PrimitivePy : public Primitive { private: py::function GetComputeFunction() const; - void ConvertCTensorToPyTensor(const py::tuple &input_args, py::tuple *convert_args) const; - void CheckHookConsistency(const py::object &grad_out, const py::object &expected_grad_out) const; + void CheckHookConsistency(const py::object &grad_out, const py::object &expected_grad_out, const py::object &code_obj, + const py::object &co_name) const; py::object python_obj_; - PrimitivePyAdapterPtr adapter_; - py::function hook_; std::string bprop_cls_name_; + PrimitivePyAdapterPtr adapter_; std::vector signatures_; + std::vector bprop_cut_prims_; + std::map backward_hook_fn_; static std::map hook_grad_; }; @@ -96,11 +95,12 @@ class PrimitivePyAdapter { void AddPyAttr(const py::str &name, const py::object &obj); void DelPyAttr(const py::str &name); py::dict GetAttrDict(); + int AddBackwardHookFn(const py::function &backward_hook_fn); + void RemoveBackwardHookFn(int key); void set_prim_type(const PrimType t); void set_const_prim(bool is_const_prim); void set_const_input_indexes(const std::vector &const_input_indexes); void set_signatures(const std::vector &signatures); - void set_hook(const py::function &hook); void set_instance_name(const std::string &s); void set_attached_primitive(const PrimitivePyPtr &prim); PrimitivePyPtr attached_primitive() const { return attached_primitive_.lock(); } @@ -110,15 +110,16 @@ class PrimitivePyAdapter { private: friend PrimitivePy; + bool is_const_prim_{false}; + int backward_hook_fn_key_{-1}; std::string name_; + std::string instance_name_; + PrimType prim_type_{kPrimTypeBuiltIn}; PrimitivePyWeakPtr attached_primitive_; mindspore::HashMap attrs_; - PrimType prim_type_{kPrimTypeBuiltIn}; - bool is_const_prim_{false}; std::vector const_input_indexes_; std::vector signatures_; - py::function hook_; - std::string instance_name_; + std::map backward_hook_fn_; }; } // namespace mindspore #endif // MINDSPORE_CCSRC_UTILS_PRIMITIVE_PY_H_ diff --git a/mindspore/core/base/core_ops.h b/mindspore/core/base/core_ops.h index db6c303d75f..1d11e894ecf 100644 --- a/mindspore/core/base/core_ops.h +++ b/mindspore/core/base/core_ops.h @@ -768,6 +768,7 @@ GVAR_DEF(PrimitivePtr, kPrimUpdateState, std::make_shared("UpdateStat GVAR_DEF(PrimitivePtr, kPrimPartial, std::make_shared("Partial", kSideEffectPropagate)); GVAR_DEF(PrimitivePtr, kPrimIdentity, std::make_shared("identity", kSideEffectPropagate)); GVAR_DEF(PrimitivePtr, kPrimHookBackward, std::make_shared("HookBackward")); +GVAR_DEF(PrimitivePtr, kPrimCellBackwardHook, std::make_shared("CellBackwardHook")); GVAR_DEF(PrimitivePtr, kPrimPrintShapeType, std::make_shared("PrintShapeType")); GVAR_DEF(PrimitivePtr, kPrimSameTypeShape, std::make_shared("SameTypeShape")); GVAR_DEF(PrimitivePtr, kPrimPrint, std::make_shared("Print")); diff --git a/mindspore/python/mindspore/_extends/parse/parser.py b/mindspore/python/mindspore/_extends/parse/parser.py index 27676a175b5..4db091118dd 100644 --- a/mindspore/python/mindspore/_extends/parse/parser.py +++ b/mindspore/python/mindspore/_extends/parse/parser.py @@ -123,8 +123,8 @@ def get_parse_method_of_class(obj, parse_method=None): if parse_method is not None: method_name = parse_method elif isinstance(obj, nn.Cell): - if obj.enable_hook: - method_name = "_hook_construct" + if obj.enable_backward_hook: + method_name = "run_backward_hook" else: method_name = "construct" if method_name is not None: diff --git a/mindspore/python/mindspore/common/api.py b/mindspore/python/mindspore/common/api.py index fb61fd37c93..3a37b835a93 100644 --- a/mindspore/python/mindspore/common/api.py +++ b/mindspore/python/mindspore/common/api.py @@ -505,6 +505,9 @@ class _PynativeExecutor: def is_top_cell(self): return self._executor.is_top_cell() + def set_hook_changed(self, cell): + self._executor.set_hook_changed(cell) + def __call__(self, obj, *args, **kwargs): args = args + tuple(kwargs.values()) return self._executor(obj, args) diff --git a/mindspore/python/mindspore/common/hook_handle.py b/mindspore/python/mindspore/common/hook_handle.py new file mode 100644 index 00000000000..a43533611b2 --- /dev/null +++ b/mindspore/python/mindspore/common/hook_handle.py @@ -0,0 +1,110 @@ +# Copyright 2022 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +"""The removable handle for cell hook function.""" +import weakref +from .api import _pynative_executor + + +class HookHandle: + r""" + It is the return object of Cell forward pre hook function, forward hook function and backward hook function. + It corresponds to the cell hook function and is used to remove the cell hook function by calling 'remove()'. + + Note: + It is only supported in pynative mode and works when registering or removing hook function for Cell object. + + Args: + hook_cell (Cell): The Cell object with hook function registered on. Default value: None. + hook_key (int): The key of cell hook function in dict. It is generated during cell hook function registration. + Default value: -1. + hook_type (str): The type of cell hook function: 'forward_pre_hook', 'forward_hook' or 'cell_backward_hook'. + Default value: "". + + Returns: + None. + + Supported Platforms: + ``Ascend`` ``GPU`` ``CPU`` + """ + def __init__(self, hook_cell=None, hook_key=-1, hook_type=""): + if hook_cell is not None: + self._hook_cell = weakref.ref(hook_cell) + else: + self._hook_cell = hook_cell + self._hook_key = hook_key + self._hook_type = hook_type + + def remove(self): + """ + Remove the cell hook function, which corresponds to this 'HookHandle' object. + In order to prevent running failed when switching to graph mode, it is not recommended to write in the + construct. + + Args: + None. + + Returns: + None. + + Supported Platforms: + ``Ascend`` ``GPU`` ``CPU`` + + Examples: + >>> import numpy as np + >>> import mindspore + >>> import mindspore.nn as nn + >>> from mindspore import Tensor + >>> from mindspore import context + >>> from mindspore.ops import GradOperation + >>> context.set_context(mode=context.PYNATIVE_MODE) + >>> def forward_pre_hook_fn(cell_id, inputs): + ... print("forward inputs: ", inputs) + ... + >>> class Net(nn.Cell): + ... def __init__(self): + ... super(Net, self).__init__() + ... self.mul = nn.MatMul() + ... self.handle = self.mul.register_forward_pre_hook(forward_pre_hook_fn) + ... + ... def construct(self, x, y): + ... x = x + x + ... x = self.mul(x, y) + ... return x + >>> grad = GradOperation(get_all=True) + >>> net = Net() + >>> output = grad(net)(Tensor(np.ones([1]).astype(np.float32)), Tensor(np.ones([1]).astype(np.float32))) + forward inputs: (Tensor(shape=[1], dtype=Float32, value= [ 2.00000000e+00]), Tensor(shape=[1], + dtype=Float32, value= [ 1.00000000e+00])) + >>> net.handle.remove() + >>> output = grad(net)(Tensor(np.ones([1]).astype(np.float32)), Tensor(np.ones([1]).astype(np.float32))) + >>> print(output) + (Tensor(shape=[1], dtype=Float32, value= [ 2.00000000e+00]), Tensor(shape=[1], dtype=Float32, + value= [ 2.00000000e+00])) + """ + if self._hook_cell is not None: + hook_cell = self._hook_cell() + if self._hook_type == "forward_pre_hook" and self._hook_key in hook_cell.forward_pre_hook: + del hook_cell.forward_pre_hook[self._hook_key] + _pynative_executor.set_hook_changed(hook_cell) + elif self._hook_type == "forward_hook" and self._hook_key in hook_cell.forward_hook: + del hook_cell.forward_hook[self._hook_key] + _pynative_executor.set_hook_changed(hook_cell) + elif self._hook_type == "cell_backward_hook": + hook_cell.cell_backward_hook.remove_backward_hook(self._hook_key) + + def __del__(self): + self._hook_cell = None + self._hook_key = None + self._hook_type = None diff --git a/mindspore/python/mindspore/nn/cell.py b/mindspore/python/mindspore/nn/cell.py index 8715c14fc11..d1217f9574b 100755 --- a/mindspore/python/mindspore/nn/cell.py +++ b/mindspore/python/mindspore/nn/cell.py @@ -18,12 +18,13 @@ import inspect import os import time from collections import OrderedDict - +from types import FunctionType, MethodType import numpy from mindspore._checkparam import args_type_check from mindspore import log as logger from mindspore.common.parameter import PARAMETER_NAME_DEFAULT +from mindspore.common.hook_handle import HookHandle from mindspore.context import ParallelMode from .. import context from .._c_expression import init_pipeline, update_func_graph_hyper_params, Cell_, FuncGraph, MixedPrecisionType @@ -32,8 +33,9 @@ from ..common import dtype as mstype from ..common.api import _cell_graph_executor, _pynative_executor, _check_all_tensor, cells_compile_cache from ..common.parameter import Parameter, ParameterTuple from ..common.tensor import Tensor, CSRTensor -from ..ops.operations import HookBackward, Cast +from ..ops.operations import Cast from ..ops.primitive import Primitive +from ..ops.operations import _inner_ops as inner from ..parallel._tensor import _load_tensor_by_layout @@ -80,10 +82,10 @@ class Cell(Cell_): IGNORE_LIST = ['_scope', '_cell_init_args', '_auto_prefix', '_cells', '_params', '_construct_inputs_names', '_construct_inputs_num', '_create_time', '_mindspore_flags', '_parallel_inputs_run', - '_parameter_layout_dict', '_params_list', '_tensor_list', '_phase', - '_auto_parallel_mode', '_backward_hook', '_bprop_debug', '_is_run', '_param_prefix', - '_attr_synced', 'enable_hook', 'pynative', 'requires_grad', - '_auto_parallel_compile_and_run', 'cell_type'] + '_parameter_layout_dict', '_params_list', '_tensor_list', '_phase', '_auto_parallel_mode', + 'forward_pre_hook', 'forward_hook', '_enable_forward_pre_hook', '_enable_forward_hook', + '_bprop_debug', 'enable_backward_hook', 'cell_backward_hook', '_is_run', '_param_prefix', + '_attr_synced', 'pynative', 'requires_grad', '_auto_parallel_compile_and_run', 'cell_type'] def __init__(self, auto_prefix=True, flags=None): Cell_.__init__(self, self._cell_tag) @@ -123,9 +125,13 @@ class Cell(Cell_): self._parallel_inputs_run = None if flags: self.add_flags(**flags) - self._backward_hook = None - self.enable_hook = False self._bprop_debug = False + self.forward_pre_hook = OrderedDict() + self.forward_hook = OrderedDict() + self._enable_forward_pre_hook = False + self._enable_forward_hook = False + self.enable_backward_hook = False + self.cell_backward_hook = None self.cell_type = None self._auto_parallel_compile_and_run = False self.cast = Cast() @@ -374,10 +380,14 @@ class Cell(Cell_): self.parameter_broadcast_done = True def run_construct(self, cast_inputs, kwargs): - if self.enable_hook: - output = self._hook_construct(*cast_inputs) + if self._enable_forward_pre_hook: + cast_inputs = self.run_forward_pre_hook(cast_inputs) + if self.enable_backward_hook: + output = self.run_backward_hook(*cast_inputs) else: output = self.construct(*cast_inputs, **kwargs) + if self._enable_forward_hook: + output = self.run_forward_hook(cast_inputs, output) return output def _check_construct_args(self, *inputs, **kwargs): @@ -466,9 +476,9 @@ class Cell(Cell_): # Run in Graph mode. if context._get_mode() == context.GRAPH_MODE: self._check_construct_args(*args, **kwargs) - if self.enable_hook: - raise ValueError("For 'Cell', it's not support hook function in graph mode, please use " - "context.set_context to set pynative mode.") + if self._enable_forward_pre_hook or self._enable_forward_hook or self.enable_backward_hook: + logger.warning(f"For 'Cell', it's not support hook function in graph mode, please use " + f"context.set_context to set pynative mode.") out = self.compile_and_run(*args) return out @@ -1442,33 +1452,295 @@ class Cell(Cell_): self.add_flags(auto_parallel=True) self._get_construct_inputs_number_and_name() - def _hook_construct(self, *inputs): - """Hook construct method to replace original construct method when hook function enabled.""" - inputs = self._backward_hook(*inputs) - inputs = self.construct(inputs) - outputs = self._backward_hook(inputs) - return outputs - - def register_backward_hook(self, fn): + def run_forward_pre_hook(self, inputs): """ - Set the cell backward hook function. Note that this function is only supported in pynative mode. - - Note: - fn must be defined as the following code. `cell_name` is the name of registered cell. - `grad_input` is gradient passed to the cell. `grad_output` is the gradient computed and passed to the - next cell or primitive, which may be modified and returned. - hook_fn(cell_name, grad_input, grad_output) -> Tensor or None. + Running forward pre hook function registered in Cell object. Args: - fn (function): Specifies the hook function with grad as input. + inputs: The input objects of Cell object. + Returns: + - **outputs** - New input objects or none. + + Supported Platforms: + ``Ascend`` ``GPU`` ``CPU`` + """ + cell_id = self.cls_name + "(" + str(id(self)) + ")" + for fn in self.forward_pre_hook.values(): + ret = fn(cell_id, inputs) + if ret is not None: + if not isinstance(ret, tuple): + inputs = (ret,) + else: + inputs = ret + return inputs + + def register_forward_pre_hook(self, hook_fn): + """ + Register forward pre hook function for Cell object. Note that this function is only supported in pynative mode. + + Note: + 'hook_fn' must be defined as the following code. + `cell_id` is the information of registered Cell object. `inputs` is the forward input objects passed to + the Cell. The 'hook_fn' can modify the forward input objects by returning new forward input objects. + It should have the following signature: + hook_fn(cell_id, inputs) -> new input objects or none. + In order to prevent running failed when switching to graph mode, it is not recommended to write in the + construct. + + Args: + hook_fn (function): Python function. Forward pre hook function. + + Returns: + - **handle** - The handle corresponding to the 'hook_fn'. + + Raises: + TypeError: If the `hook_fn` is not a function of python. + + Supported Platforms: + ``Ascend`` ``GPU`` ``CPU`` + + Examples: + >>> import numpy as np + >>> import mindspore + >>> import mindspore.nn as nn + >>> from mindspore import Tensor + >>> from mindspore import context + >>> from mindspore.ops import GradOperation + >>> context.set_context(mode=context.PYNATIVE_MODE) + >>> def forward_pre_hook_fn(cell_id, inputs): + ... print("forward inputs: ", inputs) + ... + >>> class Net(nn.Cell): + ... def __init__(self): + ... super(Net, self).__init__() + ... self.mul = nn.MatMul() + ... self.handle = self.mul.register_forward_pre_hook(forward_pre_hook_fn) + ... + ... def construct(self, x, y): + ... x = x + x + ... x = self.mul(x, y) + ... return x + >>> grad = GradOperation(get_all=True) + >>> net = Net() + >>> output = grad(net)(Tensor(np.ones([1]).astype(np.float32)), Tensor(np.ones([1]).astype(np.float32))) + forward inputs: (Tensor(shape=[1], dtype=Float32, value= [ 2.00000000e+00]), Tensor(shape=[1], + dtype=Float32, value= [ 1.00000000e+00])) + >>> print(output) + (Tensor(shape=[1], dtype=Float32, value= [ 2.00000000e+00]), Tensor(shape=[1], dtype=Float32, + value= [ 2.00000000e+00])) """ if context.get_context("mode") != context.PYNATIVE_MODE: - logger.warning("Hook function is only supported in pynative mode, you can use context.set_context to set " - "pynative mode.") + logger.warning(f"'register_forward_pre_hook' function is only supported in pynative mode, you can use " + f"context.set_context to set pynative mode.") + return HookHandle() + + if not isinstance(hook_fn, (FunctionType, MethodType)): + raise TypeError(f"When using 'register_forward_pre_hook(hook_fn)', the type of 'hook_fn' should be python " + f"function, but got {type(hook_fn)}.") + self._enable_forward_pre_hook = True + _pynative_executor.set_hook_changed(self) + if not hasattr(self, '_forward_pre_hook_key'): + self._forward_pre_hook_key = -1 + self._forward_pre_hook_key += 1 + self.forward_pre_hook[self._forward_pre_hook_key] = hook_fn + handle = HookHandle(self, self._forward_pre_hook_key, "forward_pre_hook") + return handle + + def run_forward_hook(self, inputs, output): + """ + Running forward hook function registered in Cell object. + + Args: + inputs: The input objects of Cell object. + output: The output object of Cell object. + + Returns: + - **output** - New output object or none. + + Supported Platforms: + ``Ascend`` ``GPU`` ``CPU`` + """ + cell_id = self.cls_name + "(" + str(id(self)) + ")" + for fn in self.forward_hook.values(): + ret = fn(cell_id, inputs, output) + if ret is not None: + output = ret + return output + + def register_forward_hook(self, hook_fn): + """ + Set the cell forward hook function. Note that this function is only supported in pynative mode. + + Note: + 'hook_fn' must be defined as the following code. + `cell_id` is the information of registered Cell object. `inputs` is the forward input objects passed to + the Cell. `output` is the forward output object of the Cell. The 'hook_fn' can modify the forward output + object by returning new forward output object. + It should have the following signature: + hook_fn(cell_id, inputs, output) -> new output object or none. + In order to prevent running failed when switching to graph mode, it is not recommended to write in the + construct. + + Args: + hook_fn (function): Python function. Forward hook function. + + Returns: + - **handle** - The handle corresponding to the 'hook_fn'. + + Raises: + TypeError: If the `hook_fn` is not a function of python. + + Supported Platforms: + ``Ascend`` ``GPU`` ``CPU`` + + Examples: + >>> import numpy as np + >>> import mindspore + >>> import mindspore.nn as nn + >>> from mindspore import Tensor + >>> from mindspore import context + >>> from mindspore.ops import GradOperation + >>> context.set_context(mode=context.PYNATIVE_MODE) + >>> def forward_hook_fn(cell_id, inputs, output): + ... print("forward inputs: ", inputs) + ... print("forward output: ", output) + ... + >>> class Net(nn.Cell): + ... def __init__(self): + ... super(Net, self).__init__() + ... self.mul = nn.MatMul() + ... self.handle = self.mul.register_forward_hook(forward_hook_fn) + ... + ... def construct(self, x, y): + ... x = x + x + ... x = self.mul(x, y) + ... return x + >>> grad = GradOperation(get_all=True) + >>> net = Net() + >>> output = grad(net)(Tensor(np.ones([1]).astype(np.float32)), Tensor(np.ones([1]).astype(np.float32))) + forward inputs: (Tensor(shape=[1], dtype=Float32, value= [ 2.00000000e+00]), Tensor(shape=[1], + dtype=Float32, value= [ 1.00000000e+00])) + forward output: 2.0 + >>> print(output) + (Tensor(shape=[1], dtype=Float32, value= [ 2.00000000e+00]), Tensor(shape=[1], dtype=Float32, + value= [ 2.00000000e+00])) + """ + if context.get_context("mode") != context.PYNATIVE_MODE: + logger.warning(f"'register_forward_hook' function is only supported in pynative mode, you can use " + f"context.set_context to set pynative mode.") + return HookHandle() + + if not isinstance(hook_fn, (FunctionType, MethodType)): + raise TypeError(f"When using 'register_forward_hook(hook_fn)', the type of 'hook_fn' should be python " + f"function, but got {type(hook_fn)}.") + self._enable_forward_hook = True + _pynative_executor.set_hook_changed(self) + if not hasattr(self, '_forward_hook_key'): + self._forward_hook_key = -1 + self._forward_hook_key += 1 + self.forward_hook[self._forward_hook_key] = hook_fn + handle = HookHandle(self, self._forward_hook_key, "forward_hook") + return handle + + def run_backward_hook(self, *inputs): + """ + Backward hook construct method to replace original construct method. + + Args: + inputs: The input objects of Cell object. + + Returns: + - **outputs** - The output objects of Cell object. + + Supported Platforms: + ``Ascend`` ``GPU`` ``CPU`` + """ + inputs = self.cell_backward_hook(*inputs) + if isinstance(inputs, tuple): + inputs = self.construct(*inputs) else: - self._backward_hook = HookBackward(fn, self.cls_name + "(" + str(id(self)) + ")") - self.enable_hook = True + inputs = self.construct(inputs) + if isinstance(inputs, tuple): + outputs = self.cell_backward_hook(*inputs) + else: + outputs = self.cell_backward_hook(inputs) + return outputs + + def register_backward_hook(self, hook_fn): + """ + Register the backward hook function. Note that this function is only supported in pynative mode. + + Note: + The 'hook_fn' must be defined as the following code. + `cell_id` is the information of registered cell. `grad_input` is the gradient passed to the cell. + `grad_output` is the gradient computed and passed to the next cell or primitive, which may be modified by + returning a new output gradient. + The 'hook_fn' should have the following signature: + hook_fn(cell_id, grad_input, grad_output) -> New output gradient or none. + The 'hook_fn' is executed in the python environment. + In order to prevent running failed when switching to graph mode, it is not recommended to write in the + construct. + + Args: + hook_fn (function): Python function. Backward hook function. + + Returns: + - **handle** - The handle corresponding to the 'hook_fn'. + + Raises: + TypeError: If the `hook_fn` is not a function of python. + + Supported Platforms: + ``Ascend`` ``GPU`` ``CPU`` + + Examples: + >>> import numpy as np + >>> import mindspore + >>> import mindspore.nn as nn + >>> from mindspore import Tensor + >>> from mindspore import context + >>> from mindspore.ops import GradOperation + >>> context.set_context(mode=context.PYNATIVE_MODE) + >>> def backward_hook_fn(cell_id, grad_input, grad_output): + ... print("backward input: ", grad_input) + ... print("backward output: ", grad_output) + ... + >>> class Net(nn.Cell): + ... def __init__(self): + ... super(Net, self).__init__() + ... self.relu = nn.ReLU() + ... self.handle = self.relu.register_backward_hook(backward_hook_fn) + ... + ... def construct(self, x): + ... x = x + x + ... x = self.relu(x) + ... return x + >>> grad = GradOperation(get_all=True) + >>> net = Net() + >>> output = grad(net)(Tensor(np.ones([1]).astype(np.float32))) + backward input: (Tensor(shape=[1], dtype=Float32, value= [ 1.00000000e+00]),) + backward output: (Tensor(shape=[1], dtype=Float32, value= [ 1.00000000e+00]),) + >>> print(output) + (Tensor(shape=[1], dtype=Float32, value= [ 2.00000000e+00]),) + """ + if context.get_context("mode") != context.PYNATIVE_MODE: + logger.warning(f"'register_backward_hook' function is only supported in pynative mode, you can use " + f"context.set_context to set pynative mode.") + return HookHandle() + + if not isinstance(hook_fn, (FunctionType, MethodType)): + raise TypeError(f"When using 'register_backward_hook(hook_fn)', the type of 'hook_fn' should be python " + f"function, but got {type(hook_fn)}.") + if self.cell_backward_hook is None: + self.enable_backward_hook = True + self.cell_backward_hook = inner.CellBackwardHook(self.cls_name + "(" + str(id(self)) + ")") + backward_hook_key = self.cell_backward_hook.register_backward_hook(hook_fn) + handle = HookHandle(self, backward_hook_key, "cell_backward_hook") + else: + backward_hook_key = self.cell_backward_hook.register_backward_hook(hook_fn) + handle = HookHandle(self, backward_hook_key, "cell_backward_hook") + return handle def set_param_ps(self, recurse=True, init_in_server=False): """ diff --git a/mindspore/python/mindspore/ops/operations/_inner_ops.py b/mindspore/python/mindspore/ops/operations/_inner_ops.py index 55ed861721f..b62c45fc691 100755 --- a/mindspore/python/mindspore/ops/operations/_inner_ops.py +++ b/mindspore/python/mindspore/ops/operations/_inner_ops.py @@ -14,7 +14,7 @@ # ============================================================================ """Inner operators.""" - +from types import FunctionType, MethodType import numpy as np from mindspore.common import Tensor @@ -1702,3 +1702,119 @@ class ParallelResizeBilinear(PrimitiveWithInfer): return {'shape': output_shape, 'dtype': x_dtype, 'value': None} + + +class CellBackwardHook(PrimitiveWithInfer): + r""" + This operator is used to hook input gradient and output gradient of Cell object. + + Note: + This operator is only used in backward hook function of Cell object in pynative mode. + + Args: + cell_id (str): Used to identify which cell obj the hook function registered on. For example, 'nn.Add()' is a + cell object. + + Inputs: + - **input** - The variable to hook. + + Outputs: + - **output** - Returns `input` directly. `CellBackwardHook` does not affect the forward result. + + Supported Platforms: + ``Ascend`` ``GPU`` ``CPU`` + + Examples: + >>> import mindspore + >>> from mindspore import Tensor + >>> from mindspore import context + >>> from mindspore.ops import GradOperation + >>> from mindspore.ops.operations import _inner_ops as inner + >>> context.set_context(mode=context.PYNATIVE_MODE) + >>> def hook_fn(grad): + ... print(grad) + ... + >>> hook = inner.CellBackwardHook() + >>> hook_fn_key = hook.register_backward_hook(hook_fn) + >>> def hook_test(x, y): + ... z = x * y + ... z = hook(z) + ... z = z * y + ... return z + ... + >>> grad_all = GradOperation(get_all=True) + >>> def backward(x, y): + ... return grad_all(hook_test)(x, y) + ... + >>> output = backward(Tensor(1, mindspore.float32), Tensor(2, mindspore.float32)) + (Tensor(shape=[], dtype=Float32, value= 2),) + >>> print(output) + (Tensor(shape=[], dtype=Float32, value= 4), Tensor(shape=[], dtype=Float32, value= 4)) + >>> hook.remove_backward_hook(hook_fn_key) + >>> output = backward(Tensor(1, mindspore.float32), Tensor(2, mindspore.float32)) + >>> print(output) + (Tensor(shape=[], dtype=Float32, value= 4), Tensor(shape=[], dtype=Float32, value= 4)) + """ + + def __init__(self, cell_id=""): + """Initialize CellBackwardHook""" + super(CellBackwardHook, self).__init__(self.__class__.__name__) + self.cell_id = cell_id + self.add_prim_attr("cell_id", cell_id) + self.init_attrs["cell_id"] = cell_id + + def infer_shape(self, *inputs_shape): + if len(inputs_shape) == 1: + return inputs_shape[0] + return inputs_shape + + def infer_dtype(self, *inputs_type): + if len(inputs_type) == 1: + return inputs_type[0] + return inputs_type + + def register_backward_hook(self, hook_fn): + r""" + This function is used to register backward hook function. Note that this function is only supported in pynative + mode. + + Note: + The 'hook_fn' must be defined as the following code. + `cell_id` is the information of registered cell. `grad_input` is the gradient passed to the cell. + `grad_output` is the gradient computed and passed to the next cell or primitive, which may be modified by + returning a new output gradient. + The 'hook_fn' should have the following signature: + hook_fn(cell_id, grad_input, grad_output) -> New output gradient or none. + The 'hook_fn' is executed in the python environment. + + Args: + hook_fn (Function): Python function. Backward hook function. + + Returns: + - **key** (int) - The key of 'hook_fn'. + + Raises: + TypeError: If the `hook_fn` is not a function of python. + """ + if not isinstance(hook_fn, (FunctionType, MethodType)): + raise TypeError(f"When using 'register_backward_hook(hook_fn)', the type of 'hook_fn' should be python " + f"function, but got {type(hook_fn)}.") + key = self.add_backward_hook_fn(hook_fn) + return key + + def remove_backward_hook(self, key): + r""" + This function is used to remove backward hook function. Note that this operation is only supported in pynative + mode. + + Note: + The 'key' is the object returned by 'register_backward_hook' function of the same CellBackwardHook + operator. + + Args: + key (int): The key corresponding to the 'hook_fn'. + + Returns: + None. + """ + self.remove_backward_hook_fn(key) diff --git a/mindspore/python/mindspore/ops/operations/debug_ops.py b/mindspore/python/mindspore/ops/operations/debug_ops.py index f0ea7cba47a..cdd35796599 100644 --- a/mindspore/python/mindspore/ops/operations/debug_ops.py +++ b/mindspore/python/mindspore/ops/operations/debug_ops.py @@ -324,28 +324,25 @@ class InsertGradientOf(PrimitiveWithInfer): class HookBackward(PrimitiveWithInfer): """ This operation is used as a tag to hook gradient in intermediate variables. Note that this function - is only supported in Pynative Mode. + is only supported in pynative mode. Note: - The hook function must be defined like `hook_fn(grad) -> Tensor or None`, - where grad is the gradient passed to the primitive and gradient may be - modified and passed to next primitive. The difference between a hook function and - callback of InsertGradientOf is that a hook function is executed in the python - environment while callback will be parsed and added to the graph. + The hook function must be defined like `hook_fn(grad) -> new gradient or None`, where the 'grad' is the + gradient passed to the primitive. The 'grad' may be modified by returning a new gradient and passed to next + primitive. The difference between a hook function and callback of InsertGradientOf is that the hook function is + executed in the python environment while callback will be parsed and added to the graph. Args: hook_fn (Function): Python function. hook function. - cell_id (str): Used to identify whether the function registered by the hook is actually registered on - the specified Cell. Where the Cell is an object. For example, 'nn.Add' is a Cell object. - The default value of cell_id is empty string(""), in this case, the system will automatically - register when registering. Add a value of cell_id, - the value of cell_id currently does not support custom values. Inputs: - - **inputs** (Tensor) - The variable to hook. + - **input** (Tensor) - The variable to hook. + + Outputs: + - **output** (Tensor) - Returns `input` directly. `HookBackward` does not affect the forward result. Raises: - TypeError: If `inputs` are not a Tensor. + TypeError: If `input` is not a tensor. TypeError: If `hook_fn` is not a function of python. Supported Platforms: @@ -353,15 +350,14 @@ class HookBackward(PrimitiveWithInfer): Examples: >>> import mindspore - >>> from mindspore import context - >>> from mindspore import Tensor >>> from mindspore import ops + >>> from mindspore import Tensor + >>> from mindspore import context >>> from mindspore.ops import GradOperation - >>> context.set_context(mode=context.PYNATIVE_MODE, device_target="GPU") - >>> def hook_fn(grad_out): - ... print(grad_out) + >>> context.set_context(mode=context.PYNATIVE_MODE) + >>> def hook_fn(grad): + ... print(grad) ... - >>> grad_all = GradOperation(get_all=True) >>> hook = ops.HookBackward(hook_fn) >>> def hook_test(x, y): ... z = x * y @@ -369,6 +365,7 @@ class HookBackward(PrimitiveWithInfer): ... z = z * y ... return z ... + >>> grad_all = GradOperation(get_all=True) >>> def backward(x, y): ... return grad_all(hook_test)(x, y) ... @@ -378,16 +375,16 @@ class HookBackward(PrimitiveWithInfer): (Tensor(shape=[], dtype=Float32, value= 4), Tensor(shape=[], dtype=Float32, value= 4)) """ - def __init__(self, hook_fn, cell_id=""): + def __init__(self, hook_fn): """Initialize HookBackward.""" super(HookBackward, self).__init__(self.__class__.__name__) - self.add_prim_attr("cell_id", cell_id) - self.init_attrs["cell_id"] = cell_id if not isinstance(hook_fn, (FunctionType, MethodType)): raise TypeError(f"For '{self.name}', the type of 'hook_fn' should be python function, " f"but got {type(hook_fn)}.") - self.register_hook(hook_fn) - self.cell_id = cell_id + self.add_prim_attr("cell_id", "") + self.init_attrs["cell_id"] = "" + self.cell_id = "" + self.add_backward_hook_fn(hook_fn) def infer_shape(self, *inputs_shape): if len(inputs_shape) == 1: