1/* Copyright 2017 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// 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 26namespace tensorflow { 27namespace grappler { 28 29// Make a program with specified number of stages and "width" ops per stage. 30namespace { 31GraphDef 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 89TrivialTestGraphInputYielder::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 98bool 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