1 /* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
2
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15
16 // The builtin inputs provide a mechanism to generate simple TensorFlow graphs
17 // and feed them as inputs to Grappler. This can be used for quick experiments
18 // or to derive small regression tests.
19
20 #include "tensorflow/cc/ops/standard_ops.h"
21
22 #include "tensorflow/core/framework/graph.pb.h"
23 #include "tensorflow/core/grappler/grappler_item.h"
24 #include "tensorflow/core/grappler/inputs/trivial_test_graph_input_yielder.h"
25
26 namespace tensorflow {
27 namespace grappler {
28
29 // Make a program with specified number of stages and "width" ops per stage.
30 namespace {
CreateGraphDef(int num_stages,int width,int tensor_size,bool use_multiple_devices,bool insert_queue,const std::vector<string> & device_names)31 GraphDef CreateGraphDef(int num_stages, int width, int tensor_size,
32 bool use_multiple_devices, bool insert_queue,
33 const std::vector<string>& device_names) {
34 using namespace ::tensorflow::ops; // NOLINT(build/namespaces)
35
36 tensorflow::Scope s = tensorflow::Scope::NewRootScope();
37
38 // x is from the feed.
39 const int batch_size = tensor_size < 0 ? 1 : tensor_size;
40 Output x = RandomNormal(s.WithOpName("x").WithDevice("/CPU:0"),
41 {batch_size, 1}, DataType::DT_FLOAT);
42
43 // Create stages.
44 std::vector<Output> last_stage;
45 last_stage.push_back(x);
46 for (int i = 0; i < num_stages; i++) {
47 std::vector<Output> this_stage;
48 for (int j = 0; j < width; j++) {
49 if (last_stage.size() == 1) {
50 Output unary_op = Square(
51 s.WithDevice(
52 device_names[use_multiple_devices ? j % device_names.size()
53 : 0]),
54 last_stage[0]);
55 this_stage.push_back(unary_op);
56 } else {
57 Output combine =
58 AddN(s.WithDevice(
59 device_names[use_multiple_devices ? j % device_names.size()
60 : 0]),
61 last_stage);
62 this_stage.push_back(combine);
63 }
64 }
65 last_stage = this_stage;
66 }
67
68 if (insert_queue) {
69 FIFOQueue queue(s.WithOpName("queue").WithDevice("/CPU:0"),
70 {DataType::DT_FLOAT});
71 QueueEnqueue enqueue(s.WithOpName("enqueue").WithDevice("/CPU:0"), queue,
72 last_stage);
73 QueueDequeue dequeue(s.WithOpName("dequeue").WithDevice("/CPU:0"), queue,
74 {DataType::DT_FLOAT});
75 QueueClose cancel(s.WithOpName("cancel").WithDevice("/CPU:0"), queue,
76 QueueClose::CancelPendingEnqueues(true));
77 last_stage = {dequeue[0]};
78 }
79
80 // Create output.
81 AddN output(s.WithOpName("y").WithDevice("/CPU:0"), last_stage);
82
83 GraphDef def;
84 TF_CHECK_OK(s.ToGraphDef(&def));
85 return def;
86 }
87 } // namespace
88
TrivialTestGraphInputYielder(int num_stages,int width,int tensor_size,bool insert_queue,const std::vector<string> & device_names)89 TrivialTestGraphInputYielder::TrivialTestGraphInputYielder(
90 int num_stages, int width, int tensor_size, bool insert_queue,
91 const std::vector<string>& device_names)
92 : num_stages_(num_stages),
93 width_(width),
94 tensor_size_(tensor_size),
95 insert_queue_(insert_queue),
96 device_names_(device_names) {}
97
NextItem(GrapplerItem * item)98 bool TrivialTestGraphInputYielder::NextItem(GrapplerItem* item) {
99 GrapplerItem r;
100 r.id = strings::StrCat("ns:", num_stages_, "/", // wrap
101 "w:", width_, "/", // wrap
102 "ts:", tensor_size_);
103 r.graph = CreateGraphDef(num_stages_, width_, tensor_size_,
104 true /*use_multiple_devices*/, insert_queue_,
105 device_names_);
106 // If the batch size is variable, we need to choose a value to create a feed
107 const int batch_size = tensor_size_ < 0 ? 1 : tensor_size_;
108 Tensor x(DT_FLOAT, TensorShape({batch_size, 1}));
109 r.feed.push_back(std::make_pair("x", x));
110 r.fetch.push_back("y");
111
112 if (insert_queue_) {
113 QueueRunnerDef queue_runner;
114 queue_runner.set_queue_name("queue");
115 queue_runner.set_cancel_op_name("cancel");
116 *queue_runner.add_enqueue_op_name() = "enqueue";
117 r.queue_runners.push_back(queue_runner);
118 }
119
120 *item = std::move(r);
121 return true;
122 }
123
124 } // end namespace grappler
125 } // end namespace tensorflow
126