diff --git a/mindspore/ccsrc/pipeline/pynative/pynative_execute.cc b/mindspore/ccsrc/pipeline/pynative/pynative_execute.cc index 65c698c891..fd4810a16d 100644 --- a/mindspore/ccsrc/pipeline/pynative/pynative_execute.cc +++ b/mindspore/ccsrc/pipeline/pynative/pynative_execute.cc @@ -1337,6 +1337,7 @@ TopCellInfoPtr GradExecutor::GetTopCell(const std::string &cell_id) const { } void GradExecutor::EnableOpGraphCache(bool is_enable) { + MS_LOG(DEBUG) << "Op cache is enable: " << is_enable; enable_op_cache_ = is_enable; const auto inst = MsContext::GetInstance(); MS_EXCEPTION_IF_NULL(inst); @@ -2041,6 +2042,7 @@ void GradExecutor::NewGraphInner(py::object *ret, const py::object &cell, const ResetTopCellInfo(pre_top_cell, args); PushHighOrderGraphStack(pre_top_cell); set_top_cell(pre_top_cell); + grad_order_ = pre_top_cell->grad_order(); return; } } else if ((top_cell()->IsSubCell(cell_id) || GetHighOrderStackSize() >= 1) && @@ -2096,7 +2098,7 @@ void GradExecutor::MakeNewTopGraph(const string &cell_id, const py::args &args, top_cell_list_.emplace_back(top_cell); PushHighOrderGraphStack(top_cell); set_top_cell(top_cell); - MS_LOG(DEBUG) << "New top graph, df_builder ptr " << df_builder.get() << " resource ptr " << resource.get(); + MS_LOG(DEBUG) << "New top graph, curr_g ptr " << curr_g_.get() << " resource ptr " << resource.get(); } void GradExecutor::SetTupleArgsToGraphInfoMap(const FuncGraphPtr &g, const py::object &args, const AnfNodePtr &node, @@ -2496,8 +2498,10 @@ FuncGraphPtr GradExecutor::GetBpropGraph(const prim::GradOperationPtr &grad, con py::object GradExecutor::CheckGraph(const py::object &cell, const py::args &args) { BaseRef ret = false; - ++grad_order_; check_graph_cell_id_ = GetCellId(cell, args); + if (!(top_cell_ != nullptr && top_cell_->cell_id() == check_graph_cell_id_ && grad_order_ >= 1)) { + ++grad_order_; + } if (!grad_is_running_) { MS_LOG(DEBUG) << "Grad not running yet"; return BaseRefToPyData(ret); @@ -2524,12 +2528,13 @@ py::object PynativeExecutor::CheckAlreadyRun(const py::object &cell, const py::a const auto &cell_id = grad_executor()->GetCellId(cell, args); std::string input_args_id; for (size_t i = 0; i < args.size(); ++i) { - input_args_id = input_args_id + GetId(args[i]) + "_"; + input_args_id += GetId(args[i]) + "_"; } // Check whether need to run forward process auto top_cell = grad_executor()->GetTopCell(cell_id); if (top_cell != nullptr) { forward_run = top_cell->forward_already_run(); + auto curr_top_cell = grad_executor()->top_cell(); grad_executor()->set_top_cell(top_cell); bool input_args_changed = !top_cell->input_args_id().empty() && top_cell->input_args_id() != input_args_id; if (forward_run && input_args_changed && top_cell->is_dynamic()) { @@ -2537,6 +2542,9 @@ py::object PynativeExecutor::CheckAlreadyRun(const py::object &cell, const py::a "forward process will run again"; forward_run = false; } + if (forward_run && grad_executor()->GetHighOrderStackSize() >= 1) { + grad_executor()->PushHighOrderGraphStack(curr_top_cell); + } } MS_LOG(DEBUG) << "Graph have already ran " << forward_run << " top cell id " << cell_id; return BaseRefToPyData(forward_run); diff --git a/mindspore/ccsrc/pipeline/pynative/pynative_execute.h b/mindspore/ccsrc/pipeline/pynative/pynative_execute.h index 5d214c28a2..6648fa7720 100644 --- a/mindspore/ccsrc/pipeline/pynative/pynative_execute.h +++ b/mindspore/ccsrc/pipeline/pynative/pynative_execute.h @@ -174,6 +174,8 @@ class GradExecutor { FuncGraphPtr curr_g() const; TopCellInfoPtr top_cell() const; void CheckNeedCompileGraph(); + void PushHighOrderGraphStack(const TopCellInfoPtr &top_cell); + size_t GetHighOrderStackSize() const { return high_order_stack_.size(); } TopCellInfoPtr GetTopCell(const string &cell_id) const; void EnableOpGraphCache(bool is_enable); bool need_renormalize() const { return need_renormalize_; } @@ -209,12 +211,10 @@ class GradExecutor { // Higher derivative bool IsNestedGrad() const; void SwitchTopcell(); - size_t GetHighOrderStackSize() const { return high_order_stack_.size(); } void MakeNestedCnode(const py::object &cell, const std::string &cell_id, const py::args &forward_args, const pipeline::ResourcePtr &resource, const py::object &out); void PushCellStack(const std::string &cell_id); void PopCellStack(); - void PushHighOrderGraphStack(const TopCellInfoPtr &top_cell); TopCellInfoPtr PopHighOrderGraphStack(); // Manage information of top cell. FuncGraphPtr GetDfbuilder(const std::string &cell_id = "");