diff --git a/mindspore/ccsrc/frontend/optimizer/ad/kpynative.cc b/mindspore/ccsrc/frontend/optimizer/ad/kpynative.cc index 0d15e4b7379..7571995046b 100644 --- a/mindspore/ccsrc/frontend/optimizer/ad/kpynative.cc +++ b/mindspore/ccsrc/frontend/optimizer/ad/kpynative.cc @@ -257,6 +257,8 @@ class KPynativeCellImpl : public KPynativeCell { FuncGraphPtr BuildBpropCutFuncGraph(const PrimitivePtr &prim, const CNodePtr &cnode); // Back propagate for MakeList or MakeTuple is generated from MetaFuncGraph. FuncGraphPtr BuildMakeSequenceBprop(const PrimitivePtr &prim, const CNodePtr &cnode); + // Set return node according to grad flag + void SetReturnNodeByGradFlag(const AnfNodePtrList &weights, bool grad_inputs, bool grad_weights); }; using KPynativeCellImplPtr = std::shared_ptr; @@ -294,67 +296,40 @@ FuncGraphPtr KPynativeCellImpl::Finish(const AnfNodePtrList &weights, bool grad_ auto sens_node = BuildOnesLikeValue(tape_, last_node_adjoint_iter->second->out()); last_node_adjoint_iter->second->AccumulateDout(sens_node); } - - for (size_t i = 0; i < weights.size(); ++i) { - TraceGuard trace_guard(std::make_shared(weights[i]->debug_info())); + // Add weights parameter + for (const auto &weight : weights) { + TraceGuard trace_guard(std::make_shared(weight->debug_info())); auto p = tape_->add_parameter(); - auto input_w = weights[i]->cast(); + auto input_w = weight->cast(); MS_EXCEPTION_IF_NULL(input_w); p->set_default_param(input_w->default_param()); } // BackPropagate sensitivity; BackPropagate(); - // Return the gradient; - AnfNodePtrList node_list{NewValueNode(prim::kPrimMakeTuple)}; - if (grad_inputs) { - for (auto input : cell_inputs_) { - auto input_adjoint_iter = anfnode_to_adjoin_.find(input); - if (input_adjoint_iter == anfnode_to_adjoin_.end()) { - // If input is not used in the network, just return zeros_like() as dout; - MS_LOG(WARNING) << "Input is not used in network, input: " << input->ToString(); - auto dout = BuildZerosLikeNode(tape_, input); - node_list.push_back(dout); - } else { - node_list.push_back(input_adjoint_iter->second->RealDout()); - } - } - } - if (grad_weights) { - for (auto weight : weights) { - auto input_adjoint_iter = anfnode_to_adjoin_.find(weight); - if (input_adjoint_iter == anfnode_to_adjoin_.end()) { - // If weight is not used in the network, just return zeros_like() as dout; - MS_LOG(WARNING) << "Weight is not used in network, weight: " << weight->ToString(); - auto input_w = weight->cast(); - MS_EXCEPTION_IF_NULL(input_w); - auto default_param = input_w->default_param(); - MS_EXCEPTION_IF_NULL(default_param); - auto dout = BuildZerosLikeValue(tape_, default_param); - node_list.push_back(dout); - } else { - node_list.push_back(input_adjoint_iter->second->RealDout()); - } - } - } - auto tape_output = tape_->NewCNode(node_list); - tape_->set_output(tape_output); + SetReturnNodeByGradFlag(weights, grad_inputs, grad_weights); // Replace AnfNode with parameter of tape_; auto mng = MakeManager({tape_}, false); auto tr = mng->Transact(); const auto ¶meters = tape_->parameters(); - for (size_t i = 0; i < cell_inputs_.size(); ++i) { + auto cell_inputs_size = cell_inputs_.size(); + for (size_t i = 0; i < cell_inputs_size; ++i) { tr.Replace(cell_inputs_[i], parameters[i]); } - // (Inputs, sens, weights) - size_t weight_offset = cell_inputs_.size() + 1; + // (Inputs, sens, weights) or (Inputs, weights) + size_t weight_offset = cell_inputs_size; + if (has_sens_arg) { + weight_offset = weight_offset + 1; + } for (size_t i = 0; i < weights.size(); ++i) { tr.Replace(weights[i], parameters[weight_offset + i]); } tr.Commit(); - DumpIR("before_final_opt.ir", tape_); + if (MsContext::GetInstance()->get_param(MS_CTX_SAVE_GRAPHS_FLAG)) { + DumpIR("before_final_opt.ir", tape_); + } return tape_; } @@ -769,5 +744,65 @@ FuncGraphPtr KPynativeCellImpl::BuildMakeSequenceBprop(const PrimitivePtr &prim, bprop_func_graph_cache[key] = b; return b; } + +void KPynativeCellImpl::SetReturnNodeByGradFlag(const AnfNodePtrList &weights, bool grad_inputs, bool grad_weights) { + AnfNodePtrList grad_inputs_list{NewValueNode(prim::kPrimMakeTuple)}; + if (grad_inputs) { + for (const auto &input : cell_inputs_) { + MS_EXCEPTION_IF_NULL(input); + auto input_adjoint_iter = anfnode_to_adjoin_.find(input); + if (input_adjoint_iter == anfnode_to_adjoin_.end()) { + // If input is not used in the network, just return zeros_like() as dout; + MS_LOG(WARNING) << "Input is not used in network, input: " << input->ToString(); + auto dout = BuildZerosLikeNode(tape_, input); + grad_inputs_list.push_back(dout); + } else { + grad_inputs_list.push_back(input_adjoint_iter->second->RealDout()); + } + } + } + + AnfNodePtrList grad_weights_list{NewValueNode(prim::kPrimMakeTuple)}; + if (grad_weights) { + for (const auto &weight : weights) { + MS_EXCEPTION_IF_NULL(weight); + auto input_adjoint_iter = anfnode_to_adjoin_.find(weight); + if (input_adjoint_iter == anfnode_to_adjoin_.end()) { + // If weight is not used in the network, just return zeros_like() as dout; + MS_LOG(WARNING) << "Weight is not used in network, weight: " << weight->ToString(); + auto input_w = weight->cast(); + MS_EXCEPTION_IF_NULL(input_w); + auto default_param = input_w->default_param(); + MS_EXCEPTION_IF_NULL(default_param); + auto dout = BuildZerosLikeValue(tape_, default_param); + grad_weights_list.push_back(dout); + } else { + grad_weights_list.push_back(input_adjoint_iter->second->RealDout()); + } + } + } + + AnfNodePtr tape_output; + if (grad_inputs && grad_weights) { + tape_output = tape_->NewCNode( + {NewValueNode(prim::kPrimMakeTuple), tape_->NewCNode(grad_inputs_list), tape_->NewCNode(grad_weights_list)}); + } else if (grad_inputs) { + tape_output = tape_->NewCNode(grad_inputs_list); + } else if (grad_weights) { + tape_output = tape_->NewCNode(grad_weights_list); + } else if (cell_inputs_.empty()) { + tape_output = tape_->NewCNode(grad_inputs_list); + } else { + auto input_adjoint_iter = anfnode_to_adjoin_.find(cell_inputs_[0]); + if (input_adjoint_iter == anfnode_to_adjoin_.end()) { + // If input is not used in the network, just return zeros_like() as dout; + MS_LOG(WARNING) << "Input is not used in network, input: " << cell_inputs_[0]->ToString(); + tape_output = BuildZerosLikeNode(tape_, cell_inputs_[0]); + } else { + tape_output = input_adjoint_iter->second->RealDout(); + } + } + tape_->set_output(tape_output); +} } // namespace ad } // namespace mindspore diff --git a/mindspore/ccsrc/pipeline/jit/pass.cc b/mindspore/ccsrc/pipeline/jit/pass.cc index 9cdbf015533..db55ebcd665 100644 --- a/mindspore/ccsrc/pipeline/jit/pass.cc +++ b/mindspore/ccsrc/pipeline/jit/pass.cc @@ -75,6 +75,24 @@ bool SimplifyDataStructuresPass(const ResourcePtr &res) { return true; } +bool TransformTopGraphPass(const ResourcePtr &res) { + if (res->func_graph() == nullptr) { + MS_LOG(EXCEPTION) << "Transform top graph error."; + } + FuncGraphPtr func_graph = res->func_graph(); + if (opt::FuncGraphHasTupleInput(func_graph)) { + opt::GraphTupleParamTransform graph_trans; + func_graph = graph_trans(func_graph, res->manager()); + res->set_func_graph(func_graph); + AbstractBasePtrList abs_spec_list; + auto ¶ms = func_graph->parameters(); + std::transform(params.begin(), params.end(), std::back_inserter(abs_spec_list), + [](AnfNodePtr node) { return node->abstract(); }); + res->set_args_spec(abs_spec_list); + } + return true; +} + bool CleanAfterOptAPass(const ResourcePtr &res) { MS_EXCEPTION_IF_NULL(res->func_graph()); @@ -162,6 +180,9 @@ FuncGraphPtr PrimBpOptPassStep2(const opt::irpass::OptimizeIRPassLib &irpass, co FuncGraphPtr BpropGraphFinalOptPass(const ResourcePtr &res) { MS_EXCEPTION_IF_NULL(res); + if (!TransformTopGraphPass(res)) { + MS_LOG(EXCEPTION) << "Run TransformTopGraphPass failed"; + } opt::irpass::OptimizeIRPassLib irpass; opt::OptPassConfig bg_final_opt_ = opt::OptPassConfig({ @@ -170,7 +191,6 @@ FuncGraphPtr BpropGraphFinalOptPass(const ResourcePtr &res) { irpass.depend_value_elim_, irpass.reshape_eliminate_, }); - OptPassGroupMap map({{"ad_final_opt_", bg_final_opt_}}); auto bprop_graph_final_opt = opt::Optimizer::MakeOptimizer("bprop_graph_final_opt", res, map); @@ -178,6 +198,7 @@ FuncGraphPtr BpropGraphFinalOptPass(const ResourcePtr &res) { WITH(MsProfile::GetProfile()->Step("bprop_graph_final_opt"))[&bprop_graph_final_opt, &func_graph]() { func_graph = bprop_graph_final_opt->step(func_graph, true); }; + return func_graph; } @@ -560,24 +581,6 @@ bool CconvPass(const ResourcePtr &res) { return true; } -bool TransformTopGraphPass(const ResourcePtr &res) { - if (res->func_graph() == nullptr) { - MS_LOG(EXCEPTION) << "Transform top graph error."; - } - FuncGraphPtr func_graph = res->func_graph(); - if (opt::FuncGraphHasTupleInput(func_graph)) { - opt::GraphTupleParamTransform graph_trans; - func_graph = graph_trans(func_graph, res->manager()); - res->set_func_graph(func_graph); - AbstractBasePtrList abs_spec_list; - auto ¶ms = func_graph->parameters(); - std::transform(params.begin(), params.end(), std::back_inserter(abs_spec_list), - [](AnfNodePtr node) { return node->abstract(); }); - res->set_args_spec(abs_spec_list); - } - return true; -} - bool PipelineSplitPass(const ResourcePtr &res) { return PipelineSplit(res); } void UpdateFuncGraphParameter(const FuncGraphPtr &func_graph) { diff --git a/mindspore/ccsrc/pipeline/pynative/pynative_execute.cc b/mindspore/ccsrc/pipeline/pynative/pynative_execute.cc index b74d6064983..ad74fae99b4 100644 --- a/mindspore/ccsrc/pipeline/pynative/pynative_execute.cc +++ b/mindspore/ccsrc/pipeline/pynative/pynative_execute.cc @@ -679,18 +679,21 @@ bool TopCellInfo::IsSubCell(const std::string &cell_id) const { } void TopCellInfo::clear() { + MS_LOG(DEBUG) << "Clear top cell info. Cell id " << cell_id_; + op_num_ = 0; is_dynamic_ = false; vm_compiled_ = false; is_init_kpynative_ = false; forward_already_run_ = false; - op_num_ = 0; input_args_id_.clear(); all_op_info_.clear(); - k_pynative_cell_ptr_ = nullptr; + if (resource_ != nullptr) { + resource_->Clean(); + resource_ = nullptr; + } df_builder_ = nullptr; - resource_->Clean(); - resource_ = nullptr; + k_pynative_cell_ptr_ = nullptr; graph_info_map_.clear(); sub_cell_list_.clear(); op_info_with_tensor_id_.clear(); @@ -1308,7 +1311,8 @@ void GradExecutor::UpdateForwardTensorInfoInBpropGraph(const OpExecInfoPtr &op_e } MS_EXCEPTION_IF_NULL(top_cell_); MS_EXCEPTION_IF_NULL(op_exec_info); - MS_LOG(DEBUG) << "Current op info: " << op_exec_info->op_info; + auto op_info = op_exec_info->op_info; + MS_LOG(DEBUG) << "Current op info: " << op_info; // Get input tensors; std::vector all_op_tensors; for (size_t i = 0; i < op_exec_info->op_inputs.size(); ++i) { @@ -1318,7 +1322,7 @@ void GradExecutor::UpdateForwardTensorInfoInBpropGraph(const OpExecInfoPtr &op_e TensorValueToTensor(parse::data_converter::PyDataToValue(out_real), &all_op_tensors); // Save all tensors info of current op if (need_construct_graph()) { - SaveOpInfo(top_cell_, op_exec_info->op_info, all_op_tensors); + SaveOpInfo(top_cell_, op_info, all_op_tensors); } if (already_run_top_cell_.find(top_cell_->cell_id()) == already_run_top_cell_.end()) { @@ -1332,7 +1336,11 @@ void GradExecutor::UpdateForwardTensorInfoInBpropGraph(const OpExecInfoPtr &op_e // Non-first run const auto &pre_top_cell = already_run_top_cell_.at(top_cell_->cell_id()); MS_EXCEPTION_IF_NULL(pre_top_cell); - const auto &pre_op_tensor_id = pre_top_cell->op_info_with_tensor_id().at(op_exec_info->op_info); + if (pre_top_cell->op_info_with_tensor_id().find(op_info) == pre_top_cell->op_info_with_tensor_id().end()) { + MS_LOG(EXCEPTION) << "Can not find op info " << op_info << " in op info with tensor id map. Top cell " + << top_cell_->cell_id(); + } + const auto &pre_op_tensor_id = pre_top_cell->op_info_with_tensor_id().at(op_info); if (pre_op_tensor_id.size() != all_op_tensors.size()) { MS_LOG(EXCEPTION) << "The size of pre op tensor id: " << pre_op_tensor_id.size() << " is not equal to the size of all tensors of current op " << all_op_tensors.size(); @@ -1692,13 +1700,20 @@ void GradExecutor::ClearCellRes(const std::string &cell_id) { it->clear(); } top_cell_list_.clear(); + already_run_top_cell_.clear(); + MS_LOG(DEBUG) << "Clear all cell resources"; return; } for (auto it = top_cell_list_.begin(); it != top_cell_list_.end();) { - if (IsCellObjIdEq(cell_id, (*it)->cell_id())) { + auto top_cell_id = (*it)->cell_id(); + if (IsCellObjIdEq(cell_id, top_cell_id)) { (*it)->clear(); it = top_cell_list_.erase(it); - break; + if (already_run_top_cell_.find(top_cell_id) != already_run_top_cell_.end()) { + (void)already_run_top_cell_.erase(top_cell_id); + } + MS_LOG(ERROR) << "Clear top cell resource. Top cell id " << top_cell_id; + continue; } it++; } @@ -1898,14 +1913,19 @@ void GradExecutor::EndGraphInner(py::object *ret, const py::object &cell, const auto tuple = out.cast(); auto tuple_size = static_cast(tuple.size()); + ValuePtrList input_args; std::vector inputs; inputs.emplace_back(NewValueNode(prim::kPrimMakeTuple)); for (int64_t i = 0; i < tuple_size; i++) { inputs.emplace_back(GetInput(tuple[i], false)); + input_args.emplace_back(parse::data_converter::PyDataToValue(tuple[i])); } auto cnode = curr_g_->NewCNode(inputs); SetTupleArgsToGraphInfoMap(curr_g_, out, cnode); SetNodeMapInGraphInfoMap(curr_g_, out_id, cnode); + ValuePtr out_value = parse::data_converter::PyDataToValue(out); + ad::GradPynativeOp(top_cell()->k_pynative_cell_ptr(), cnode, input_args, out_value); + MS_LOG(DEBUG) << "Tuple output node info " << cnode->DebugString(); } else { MS_LOG(DEBUG) << "Set ValueNode as output for graph, out id: " << out_id; MakeValueNode(out, out_id); @@ -2017,7 +2037,7 @@ void GradExecutor::GradNetInner(py::object *ret, const GradOperationPtr &grad, c MS_LOG(DEBUG) << "df_builder ptr " << df_builder.get() << " resource ptr " << resource.get(); // Get params(weights) require derivative - auto w_args = GetWeightsArgs(weights, df_builder); + auto w_args = GetWeightsArgs(grad, weights, df_builder); // Get bprop graph of top cell auto bprop_graph = GetBpropGraph(grad, w_args, size, args); resource->set_func_graph(bprop_graph); @@ -2037,11 +2057,18 @@ void GradExecutor::GradNetInner(py::object *ret, const GradOperationPtr &grad, c resource->Clean(); } -std::vector GradExecutor::GetWeightsArgs(const py::object &weights, const FuncGraphPtr &df_builder) { - if (!py::hasattr(weights, "__parameter_tuple__")) { +std::vector GradExecutor::GetWeightsArgs(const GradOperationPtr &grad, const py::object &weights, + const FuncGraphPtr &df_builder) { + MS_EXCEPTION_IF_NULL(grad); + MS_EXCEPTION_IF_NULL(df_builder); + if (!grad->get_by_list_ && py::isinstance(weights)) { + MS_LOG(DEBUG) << "The input weight is None when run GradNetInner. Return parameters of df_builder directly"; + return df_builder->parameters(); + } else if (!py::hasattr(weights, "__parameter_tuple__")) { MS_LOG(DEBUG) << "No paramter_tuple get"; return {}; } + auto tuple = weights.cast(); MS_LOG(DEBUG) << "Get weights tuple size " << tuple.size(); std::vector w_args; @@ -2111,8 +2138,8 @@ FuncGraphPtr GradExecutor::GetBpropGraph(const GradOperationPtr &grad, const std MS_EXCEPTION_IF_NULL(grad); auto k_pynative_cell_ptr = top_cell()->k_pynative_cell_ptr(); MS_EXCEPTION_IF_NULL(k_pynative_cell_ptr); - auto bprop_graph = ad::GradPynativeCellEnd(k_pynative_cell_ptr, weights, grad->get_all_, grad->get_by_list_, - grad->sens_param_); + auto bprop_graph = + ad::GradPynativeCellEnd(k_pynative_cell_ptr, weights, grad->get_all_, grad->get_by_list_, grad->sens_param_); MS_EXCEPTION_IF_NULL(bprop_graph); MS_LOG(DEBUG) << "Top graph input params size " << arg_size; @@ -2130,7 +2157,10 @@ FuncGraphPtr GradExecutor::GetBpropGraph(const GradOperationPtr &grad, const std MS_EXCEPTION_IF_NULL(manager); manager->AddFuncGraph(bprop_graph); auto optimized_bg = pipeline::PrimBpropOptimizer::GetPrimBpropOptimizerInst().BpropGraphFinalOpt(resource); - DumpIR("after_final_opt.ir", optimized_bg); + + if (MsContext::GetInstance()->get_param(MS_CTX_SAVE_GRAPHS_FLAG)) { + DumpIR("after_final_opt.ir", optimized_bg); + } optimized_bg->ClearAllManagerInfo(); return optimized_bg; } diff --git a/mindspore/ccsrc/pipeline/pynative/pynative_execute.h b/mindspore/ccsrc/pipeline/pynative/pynative_execute.h index ae3c86ce90f..7f44f493749 100644 --- a/mindspore/ccsrc/pipeline/pynative/pynative_execute.h +++ b/mindspore/ccsrc/pipeline/pynative/pynative_execute.h @@ -175,7 +175,7 @@ class GradExecutor { bool grad_flag() const { return grad_flag_; } void set_grad_flag(bool flag) { grad_flag_ = flag; } bool in_grad_process() const { return in_grad_process_; } - bool in_cell_with_custom_bprop_() const {return custom_bprop_cell_count_ > 0;} + bool in_cell_with_custom_bprop_() const { return custom_bprop_cell_count_ > 0; } AnfNodePtr GetInput(const py::object &obj, bool op_mask); std::string GetCellId(const py::object &obj, const py::args &args); std::stack &cell_stack() { return cell_stack_; } @@ -230,7 +230,8 @@ class GradExecutor { const py::args &args); FuncGraphPtr GetBpropGraph(const GradOperationPtr &grad, const std::vector &weights, size_t arg_size, const py::args &args); - std::vector GetWeightsArgs(const py::object &weights, const FuncGraphPtr &df_builder); + std::vector GetWeightsArgs(const GradOperationPtr &grad, const py::object &weights, + const FuncGraphPtr &df_builder); abstract::AbstractBasePtrList GetArgsSpec(const py::args &args, const FuncGraphPtr &bprop_graph); void SetTupleItemArgsToGraphInfoMap(const FuncGraphPtr &g, const py::object &id, const AnfNodePtr &node, const std::vector &index_sequence, bool is_param = false);