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 #ifndef TENSORFLOW_GRAPPLER_COSTS_MEASURING_COST_ESTIMATOR_H_ 17 #define TENSORFLOW_GRAPPLER_COSTS_MEASURING_COST_ESTIMATOR_H_ 18 19 #include <string> 20 #include <utility> 21 #include <vector> 22 23 #include "tensorflow/core/framework/tensor.h" 24 #include "tensorflow/core/grappler/costs/cost_estimator.h" 25 #include "tensorflow/core/lib/core/status.h" 26 #include "tensorflow/core/lib/core/threadpool.h" 27 #include "tensorflow/core/platform/types.h" 28 29 namespace tensorflow { 30 class CostGraphDef; 31 class GraphDef; 32 } // namespace tensorflow 33 34 namespace tensorflow { 35 namespace grappler { 36 37 class Cluster; 38 struct GrapplerItem; 39 40 // Estimate the cost of running a Grappler item by actually running the 41 // corresponding TensorFlow graph on the specified cluster and measuring the 42 // runtimes. 43 class MeasuringCostEstimator : public CostEstimator { 44 public: 45 // Run the model for measurement_steps to measure its average cost. 46 // When measurement_threads is greater than 0, use a threadpool of as many 47 // threads to run the measurements; otherwise, run them serially. Does not 48 // take ownership of cluster. 49 explicit MeasuringCostEstimator(Cluster* cluster, int measurement_steps, 50 int measurement_threads); 51 ~MeasuringCostEstimator() override {} 52 53 // Initializes the estimator for the specified grappler item. 54 // This implementation always returns OK. 55 Status Initialize(const GrapplerItem& item) override; 56 57 // Runs the optimized version of the graph on the cluster, measure 58 // the runtimes of each operation, and annotated the CostGraphDef 59 // with the corresponding measurements. 60 // Returns the average latency for the whole graph. 61 Status PredictCosts(const GraphDef& optimized_graph, CostGraphDef* cost_graph, 62 Costs* overall_cost) const override; 63 64 private: 65 Cluster* cluster_; // Not owned. 66 int measurement_steps_; 67 int measurement_threads_; 68 std::vector<std::pair<string, Tensor>> feed_; 69 std::vector<string> fetch_; 70 std::unique_ptr<thread::ThreadPool> thread_pool_; 71 }; 72 73 } // end namespace grappler 74 } // end namespace tensorflow 75 76 #endif // TENSORFLOW_GRAPPLER_COSTS_MEASURING_COST_ESTIMATOR_H_ 77
