1/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
2
3Licensed under the Apache License, Version 2.0 (the "License");
4you may not use this file except in compliance with the License.
5You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9Unless required by applicable law or agreed to in writing, software
10distributed under the License is distributed on an "AS IS" BASIS,
11WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12See the License for the specific language governing permissions and
13limitations under the License.
14==============================================================================*/
15
16#ifndef TENSORFLOW_GRAPH_NODE_BUILDER_H_
17#define TENSORFLOW_GRAPH_NODE_BUILDER_H_
18
19#include <vector>
20#include "tensorflow/core/framework/node_def_builder.h"
21#include "tensorflow/core/framework/op.h"
22#include "tensorflow/core/framework/op_def.pb.h"
23#include "tensorflow/core/graph/graph.h"
24#include "tensorflow/core/lib/core/status.h"
25#include "tensorflow/core/lib/core/stringpiece.h"
26#include "tensorflow/core/lib/gtl/array_slice.h"
27
28namespace tensorflow {
29
30// This is a helper for creating a Node and adding it to a Graph.
31// Internally, it uses a NodeDefBuilder to automatically set attrs
32// that can be inferred from the inputs, and use default values
33// (where they exist) for unspecified attrs. Example usage:
34//
35// Node* node;
36// Status status = NodeBuilder(node_name, op_name)
37// .Input(...)
38// .Attr(...)
39// .Finalize(&graph, &node);
40// if (!status.ok()) return status;
41// // Use node here.
42class NodeBuilder {
43 public:
44 // For specifying the output of a Node to provide to one of the Input()
45 // functions below. It supports both regular inputs (where you are
46 // connecting to an existing Node*), and inputs from outside the graph
47 // (or haven't been added to the graph yet, like back edges, where
48 // you don't have a Node*). Both types can be mixed, e.g. in an
49 // ArraySlice.
50 struct NodeOut {
51 // For referencing an existing Node.
52 NodeOut(Node* n, int32 i = 0);
53
54 // For referencing Nodes not in the graph being built. It is
55 // useful when preparing a graph for ExtendSession or creating a
56 // back edge to a node that hasn't been added to the graph yet,
57 // but will be.
58 NodeOut(StringPiece name, int32 i, DataType t);
59
60 // Default constructor for std::vector<NodeOut>.
61 NodeOut();
62
63 Node* node;
64 // error is set to true if:
65 // * the NodeOut was default constructed and never overwritten,
66 // * a nullptr Node* was passed to the NodeOut constructor, or
67 // * an out-of-range index was passed to the NodeOut constructor.
68 bool error;
69 string name;
70 int32 index;
71 DataType dt;
72 };
73
74 // Specify the name and the Op (either via an OpDef or the name of
75 // the Op plus a registry) for the Node. Other fields are
76 // specified by calling the methods below.
77 // REQUIRES: The OpDef must satisfy ValidateOpDef().
78 NodeBuilder(StringPiece name, StringPiece op_name,
79 const OpRegistryInterface* op_registry = OpRegistry::Global());
80 NodeBuilder(StringPiece name, const OpDef* op_def);
81
82 // Create a NodeBuilder from an existing NodeDefBuilder.
83 NodeBuilder(const NodeDefBuilder& def_builder);
84
85 // You must call one Input() function per input_arg in the Op,
86 // *and in the same order as the input_args appear in the OpDef.*
87
88 // For inputs that take a single tensor.
89 NodeBuilder& Input(Node* src_node, int src_index = 0);
90 NodeBuilder& Input(NodeOut src);
91
92 // For inputs that take a list of tensors.
93 NodeBuilder& Input(gtl::ArraySlice<NodeOut> src_list);
94
95 // Require that this node run after src_node(s).
96 NodeBuilder& ControlInput(Node* src_node);
97 NodeBuilder& ControlInputs(gtl::ArraySlice<Node*> src_nodes);
98
99 // Sets the "requested device spec" in the NodeDef (not the
100 // "assigned device" in the Node).
101 NodeBuilder& Device(StringPiece device_spec);
102
103 // Set the value of an attr. attr_name must match the name of one of
104 // attrs defined by the Op, and value must have the corresponding type
105 // (see SetAttrValue() in ../framework/attr_value_util.h for legal
106 // types for value). Note that attrs will be set automatically if
107 // they can be determined by the inputs.
108 template <class T>
109 NodeBuilder& Attr(StringPiece attr_name, T&& value);
110 template <class T>
111 NodeBuilder& Attr(StringPiece attr_name, std::initializer_list<T> value);
112
113 // Validates the described node and adds it to *graph, adding edges
114 // for all (non-back) inputs. If created_node is not nullptr,
115 // *created_node will be set to the new node (or nullptr on error).
116 Status Finalize(Graph* graph, Node** created_node) const;
117
118 // Accessors for the values set in the constructor.
119 const string& node_name() const { return def_builder_.node_name(); }
120 const OpDef& op_def() const { return def_builder_.op_def(); }
121
122 private:
123 static DataType SafeGetOutput(const Node* node, int i, bool* error) {
124 if (node != nullptr && i >= 0 && i < node->num_outputs()) {
125 *error = false;
126 return node->output_type(i);
127 } else {
128 *error = true;
129 return DT_FLOAT;
130 }
131 }
132
133 // If SafeGetOutput indicates a range error, add it to errors_.
134 void AddIndexError(const Node* node, int i);
135
136 // Set *dt and returns true if i is in range. Combines
137 // SafeGetOutput() and AddIndexError().
138 bool GetOutputType(const Node* node, int i, DataType* dt);
139
140 NodeDefBuilder def_builder_;
141 std::vector<NodeOut> inputs_;
142 std::vector<Node*> control_inputs_;
143 std::vector<string> errors_;
144};
145
146// IMPLEMENTATION -------------------------------------------------------------
147
148template <class T>
149NodeBuilder& NodeBuilder::Attr(StringPiece attr_name, T&& value) {
150 def_builder_.Attr(attr_name, std::forward<T>(value));
151 return *this;
152}
153
154template <class T>
155NodeBuilder& NodeBuilder::Attr(StringPiece attr_name,
156 std::initializer_list<T> value) {
157 def_builder_.Attr(attr_name, value);
158 return *this;
159}
160
161} // namespace tensorflow
162
163#endif // TENSORFLOW_GRAPH_NODE_BUILDER_H_
164