diff --git a/ge/graph/passes/net_output_pass.cc b/ge/graph/passes/net_output_pass.cc index e3f2b71a..c6ab062a 100644 --- a/ge/graph/passes/net_output_pass.cc +++ b/ge/graph/passes/net_output_pass.cc @@ -103,6 +103,12 @@ Status NetOutputPass::GetOutputNode(const ge::ComputeGraphPtr &graph, std::vecto GELOGI("user set out node [%s] is found in user def targets, out node is prio!", ele.first->GetName().c_str()); targets_.erase(iter); } + + auto op_desc = ele.first->GetOpDesc(); + GE_CHECK_NOTNULL(op_desc); + if (op_desc->HasAttr(ATTR_ATC_USER_DEFINE_OUTPUT_NODES)) { + is_user_define_ouput_nodes = true; + } output_nodes_info.push_back({ele.first, ele.second, -1}); } GELOGI("Output node set by user or leaf node, size:%zu.", output_nodes_info.size()); @@ -414,7 +420,7 @@ Status NetOutputPass::ProcessWithNetoutput(const ge::ComputeGraphPtr &graph, con Status NetOutputPass::AddCtrlEdgesBetweenLeafAndNetOutput(const ge::ComputeGraphPtr &graph, const ge::NodePtr &net_out_node) { GE_CHECK_NOTNULL(net_out_node); - if (!GetLocalOmgContext().user_out_nodes.empty()) { + if (!GetLocalOmgContext().user_out_nodes.empty() || is_user_define_ouput_nodes) { GELOGI("No need to add ctrl edge to netoutput because user out nodes have been set."); return SUCCESS; } diff --git a/ge/graph/passes/net_output_pass.h b/ge/graph/passes/net_output_pass.h index b959bd96..ab190169 100644 --- a/ge/graph/passes/net_output_pass.h +++ b/ge/graph/passes/net_output_pass.h @@ -220,6 +220,7 @@ class NetOutputPass : public GraphPass { bool is_include_special_node_ = false; std::set targets_; friend class ReUpdateNetOutputPass; + bool is_user_define_ouput_nodes = false; }; } // namespace ge #endif // GE_GRAPH_PASSES_NET_OUTPUT_PASS_H_ diff --git a/ge/session/omg.cc b/ge/session/omg.cc index 104b3d00..16449363 100755 --- a/ge/session/omg.cc +++ b/ge/session/omg.cc @@ -485,6 +485,10 @@ Status SetOutputNodeInfo(ge::Graph &graph, const std::string &output_type, const GELOGE(domi::FAILED, "Check out node (%s) fail.", user_out_nodes[i].first.c_str()); return domi::FAILED; } + + // add user_define_output_nodes attr. + (void)ge::AttrUtils::SetStr(op_desc, ATTR_ATC_USER_DEFINE_OUTPUT_NODES, "true"); + if (i < output_formats.size()) { if (output_formats[i] == domi::DOMI_TENSOR_NC1HWC0) { GELOGI("The output node [%s] should be set NC1HWC0", user_out_nodes[i].first.c_str()); diff --git a/metadef b/metadef index e2a36e47..37465b85 160000 --- a/metadef +++ b/metadef @@ -1 +1 @@ -Subproject commit e2a36e47259da923702971dcfa8a78d100b74902 +Subproject commit 37465b85d30b67a0edcc6ea4acd2f11a9697c7af diff --git a/parser b/parser index 0b7d8c9f..5fa1f3ed 160000 --- a/parser +++ b/parser @@ -1 +1 @@ -Subproject commit 0b7d8c9ffba6de83c4232db4f2e105ba51cd6296 +Subproject commit 5fa1f3ed9b1785b9fd1623d624de91838dff615e