From 8edb90f4e435178d19ed37946411fa4f1e3b106c Mon Sep 17 00:00:00 2001 From: Scott McKay Date: Mon, 18 Mar 2019 19:33:52 +1000 Subject: [PATCH 1/2] Update the GraphProto for subgraphs when saving the Graph. This is required to produce a valid overall Graph if the Graph has been optimized. --- include/onnxruntime/core/graph/graph.h | 8 ++++++-- onnxruntime/core/graph/graph.cc | 27 ++++++++++++++++++-------- 2 files changed, 25 insertions(+), 10 deletions(-) diff --git a/include/onnxruntime/core/graph/graph.h b/include/onnxruntime/core/graph/graph.h index 39dfbf3d345b6..02d40a0534a09 100644 --- a/include/onnxruntime/core/graph/graph.h +++ b/include/onnxruntime/core/graph/graph.h @@ -279,8 +279,12 @@ class Node { /** Sets the execution ProviderType that this Node will be executed by. */ void SetExecutionProviderType(ProviderType execution_provider_type); - /** Gets the NodeProto representation of this Node. */ - void ToProto(ONNX_NAMESPACE::NodeProto& proto) const; + /** Gets the NodeProto representation of this Node. + @param update_subgraphs Update the GraphProto values for subgraphs in the returned NodeProto. + If graph optimization has been run this is most likely required + to ensure the complete Graph is valid. + */ + void ToProto(ONNX_NAMESPACE::NodeProto& proto, bool update_subgraphs = false) const; /** Call the provided function for all explicit inputs, implicit inputs, and outputs of this Node. If the NodeArg is an explicit or implicit input, is_input will be true when func is called. diff --git a/onnxruntime/core/graph/graph.cc b/onnxruntime/core/graph/graph.cc index dcf394813eb7f..509a61e86e2f9 100644 --- a/onnxruntime/core/graph/graph.cc +++ b/onnxruntime/core/graph/graph.cc @@ -20,6 +20,10 @@ #include "core/common/logging/logging.h" #include "onnx/checker.h" #include "core/graph/schema_registry.h" + +// TEMP +#include "core/platform/env.h" + using namespace ONNX_NAMESPACE; using namespace ONNX_NAMESPACE::Utils; using namespace ONNX_NAMESPACE::checker; @@ -318,21 +322,26 @@ void Node::SetExecutionProviderType(ProviderType execution_provider_type) { execution_provider_type_ = execution_provider_type; } -void Node::ToProto(NodeProto& proto) const { - // Set name. +void Node::ToProto(NodeProto& proto, bool update_subgraphs) const { proto.set_name(name_); - // Set op type. proto.set_op_type(op_type_); - // Set op domain; - proto.set_domain(domain_); - // Set doc string. - proto.set_doc_string(description_); + + if (!domain_.empty()) + proto.set_domain(domain_); + + if (!description_.empty()) + proto.set_doc_string(description_); // Set attributes. proto.clear_attribute(); for (auto attribute : attributes_) { const gsl::not_null attr{proto.add_attribute()}; *attr = attribute.second; + if (update_subgraphs && attribute.second.has_g()) { + GraphProto latest = attr_to_subgraph_map_.find(attribute.first)->second->ToGraphProto(); + attr->clear_g(); + *attr->mutable_g() = latest; + } } // Set inputs' definitions. @@ -628,6 +637,8 @@ Graph::Graph(GraphProto* graph_proto, *tensor = node.attribute(0).t(); *(tensor->mutable_name()) = node.output(0); + name_to_initial_tensor_[tensor->name()] = tensor; + // we remove the node and add it as an initializer, but still need it to appear in the // graph inputs to make the ONNX checker happy. add a new input due to that. auto graph_inputs = graph_proto_->mutable_input(); @@ -2104,7 +2115,7 @@ const GraphProto& Graph::ToGraphProto() { for (auto& node_idx : graph_viewer.GetNodesInTopologicalOrder()) { const gsl::not_null node_proto{graph_proto_->add_node()}; const gsl::not_null p_node{GetNode(node_idx)}; - p_node->ToProto(*node_proto); + p_node->ToProto(*node_proto, true); } if (!removed_initializer_indexes_.empty()) { From 8208835d54fd5fbe99f2d8ae51e19a252187c99c Mon Sep 17 00:00:00 2001 From: Scott McKay Date: Mon, 18 Mar 2019 19:38:06 +1000 Subject: [PATCH 2/2] Add a couple of unsaved cleanups. --- include/onnxruntime/core/graph/graph.h | 4 ++-- onnxruntime/core/graph/graph.cc | 6 ------ 2 files changed, 2 insertions(+), 8 deletions(-) diff --git a/include/onnxruntime/core/graph/graph.h b/include/onnxruntime/core/graph/graph.h index 02d40a0534a09..bbcb8583a641a 100644 --- a/include/onnxruntime/core/graph/graph.h +++ b/include/onnxruntime/core/graph/graph.h @@ -280,8 +280,8 @@ class Node { void SetExecutionProviderType(ProviderType execution_provider_type); /** Gets the NodeProto representation of this Node. - @param update_subgraphs Update the GraphProto values for subgraphs in the returned NodeProto. - If graph optimization has been run this is most likely required + @param update_subgraphs Update the GraphProto values for any subgraphs in the returned NodeProto. + If graph optimization has been run this is most likely required to ensure the complete Graph is valid. */ void ToProto(ONNX_NAMESPACE::NodeProto& proto, bool update_subgraphs = false) const; diff --git a/onnxruntime/core/graph/graph.cc b/onnxruntime/core/graph/graph.cc index 509a61e86e2f9..4b9a775818e02 100644 --- a/onnxruntime/core/graph/graph.cc +++ b/onnxruntime/core/graph/graph.cc @@ -20,10 +20,6 @@ #include "core/common/logging/logging.h" #include "onnx/checker.h" #include "core/graph/schema_registry.h" - -// TEMP -#include "core/platform/env.h" - using namespace ONNX_NAMESPACE; using namespace ONNX_NAMESPACE::Utils; using namespace ONNX_NAMESPACE::checker; @@ -637,8 +633,6 @@ Graph::Graph(GraphProto* graph_proto, *tensor = node.attribute(0).t(); *(tensor->mutable_name()) = node.output(0); - name_to_initial_tensor_[tensor->name()] = tensor; - // we remove the node and add it as an initializer, but still need it to appear in the // graph inputs to make the ONNX checker happy. add a new input due to that. auto graph_inputs = graph_proto_->mutable_input();