Domain exists in GraphProto but not in Node (onnx#1310)

* Domain needs to be copied in graph nodes * Initialize has_domain_
JSnobody · Aug 22, 2018 · 2cbf740 · 2cbf740
1 parent 9b874e9
commit 2cbf740
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 0 deletions.
diff --git a/onnx/common/ir.h b/onnx/common/ir.h
@@ -409,6 +409,8 @@ struct Node : public Attributes<Node> {
   size_t stage_;
   bool has_name_;
   std::string name_;
+  bool has_domain_;
+  std::string domain_;
   bool has_doc_string_;
   std::string doc_string_;
 
@@ -426,6 +428,16 @@ struct Node : public Attributes<Node> {
     has_name_ = true;
     name_ = std::move(name);
   }
+  bool has_domain() {
+    return has_domain_;
+  }
+  const std::string& domain() const {
+    return domain_;
+  }
+  void setDomain(std::string domain) {
+    has_domain_ = true;
+    domain_ = std::move(domain);
+  }
   bool has_doc_string() const {
     return has_doc_string_;
   }
@@ -1132,6 +1144,7 @@ inline Node::Node(Graph * graph_, NodeKind kind_) :
   graph_(graph_),
   stage_(graph_->new_node_stage_),
   has_name_(false),
+  has_domain_(false),
   has_doc_string_(false) {
   graph_->all_nodes.emplace(this);
 }

diff --git a/onnx/common/ir_pb_converter.cc b/onnx/common/ir_pb_converter.cc
@@ -252,6 +252,9 @@ std::unique_ptr<Graph> graphProtoToGraph(const ONNX_NAMESPACE::GraphProto& gp, b
     if (np.has_name()) {
       n->setName(np.name());
     }
+    if (np.has_domain()) {
+      n->setDomain(np.domain());
+    }
   }
 
   for (auto n : g->nodes()) {
@@ -528,6 +531,9 @@ void encodeGraph(GraphProto * p_g, const std::shared_ptr<Graph> & g) {
     if (node->has_name()) {
       p_n->set_name(node->name());
     }
+    if (node->has_domain()) {
+      p_n->set_domain(node->domain());
+    }
   }
 
   auto num_initializers = g->initializers().size();