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 #ifndef TENSORFLOW_CONTRIB_LITE_TOCO_TENSORFLOW_GRAPH_MATCHING_RESOLVE_CLUSTER_H 16 #define TENSORFLOW_CONTRIB_LITE_TOCO_TENSORFLOW_GRAPH_MATCHING_RESOLVE_CLUSTER_H 17 18 #include <string> 19 #include <unordered_map> 20 #include <vector> 21 22 #include "tensorflow/contrib/lite/toco/tensorflow_graph_matching/cluster.h" 23 #include "tensorflow/contrib/lite/toco/tensorflow_graph_matching/resolve_svdf.h" 24 #include "tensorflow/core/framework/graph.pb.h" 25 #include "tensorflow/core/framework/node_def.pb.h" 26 27 namespace toco { 28 29 // Given a graph info and a list of cluster classes (cluster_factories), it 30 // partitions the graph to clusters, and then collapses each cluster into their 31 // corresponding composite ops. It generates a new graph using the newly 32 // generated composite ops. Each cluster factory is responsible to recognize a 33 // cluster of nodes into a cluster using a name-based pattern matching approach. 34 std::unique_ptr<tensorflow::GraphDef> MaybeResolveClusters( 35 const tensorflow::GraphDef& graph_def, 36 const std::vector<ClusterFactoryInterface*>& cluster_factories); 37 38 // Adds a node to a given graph. The added node will be a copy of a given source 39 // node, except for the inputs. If the inputs are coming from a node which 40 // belongs to another cluster, then those inputs are renamed to the source 41 // cluster name. 42 void AddNodeToGraph(const tensorflow::NodeDef& node, 43 const std::vector<string>& cluster_names, 44 tensorflow::GraphDef* graph); 45 46 // Given a graph and a cluster class, it finds all the nodes which belong to a 47 // given class factory, encapsulate them inside a cluster of the given type and 48 // returns a vector of those clusters. It also labels the nodes in that graph if 49 // they belong to the generated clusters. 50 bool FindCluster(const ClusterFactoryInterface& cluster_factory, 51 const tensorflow::GraphDef& graph_def, 52 std::unordered_map<string, bool>* is_node_in_cluster, 53 std::vector<std::unique_ptr<Cluster>>* clusters); 54 55 // Receives a graph and generates another graph by replacing the cluster of 56 // nodes which matches a given composite op. Each composite op is represented 57 // using a class factory. 58 std::unique_ptr<tensorflow::GraphDef> MaybeReplaceCompositeSubgraph( 59 const tensorflow::GraphDef& tf_graph); 60 61 } // end namespace toco 62 63 #endif // CONTRIB_LITE_TOCO_TENSORFLOW_GRAPH_MATCHING_RESOLVE_CLUSTER_H 64
