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_LITE_TOCO_TENSORFLOW_GRAPH_MATCHING_RESOLVE_SVDF_H_ 16 #define TENSORFLOW_LITE_TOCO_TENSORFLOW_GRAPH_MATCHING_RESOLVE_SVDF_H_ 17 18 #include <string> 19 #include <vector> 20 21 #include "tensorflow/lite/toco/model.h" 22 #include "tensorflow/lite/toco/tensorflow_graph_matching/cluster.h" 23 #include "tensorflow/lite/toco/tensorflow_graph_matching/cluster_utils.h" 24 #include "tensorflow/lite/toco/tooling_util.h" 25 #include "tensorflow/core/framework/attr_value.pb.h" 26 #include "tensorflow/core/framework/graph.pb.h" 27 #include "tensorflow/core/framework/node_def.pb.h" 28 29 namespace toco { 30 31 class SvdfCluster : public Cluster { 32 public: 33 // For this cluster, it collapses all the nodes in nodes_ into a composite op 34 // and it returns all the newly generated ops in new_nodes_. 35 void CreateNodes() override; 36 37 // A helper function to set the pattern of Const nodes which CreateNodes() 38 // should handle specially. AddConstNodePattern(const std::string & const_pattern)39 void AddConstNodePattern(const std::string& const_pattern) { 40 const_node_patterns_.push_back(const_pattern); 41 } 42 ~SvdfCluster()43 virtual ~SvdfCluster() {} 44 45 private: 46 // The main function which is used to create Const nodes for this cluster. 47 // These Const nodes are the inputs to the composite op generated for this 48 // cluster. 49 void CreateConstNode(const std::string& const_pattern); 50 51 // Receives a vector of Const nodes, merge them (if necessary) and returns 52 // only one Const node holding all the arrays contents. It transposes it if 53 // needed. 54 void MaybeMergeConstNodes( 55 const std::vector<const tensorflow::NodeDef*>& const_node_parts, 56 bool transpose_tensor_value, 57 const std::unique_ptr<tensorflow::NodeDef>& merged_node); 58 59 // Infer the value of Svdf filter rank, by looking up a reshape operator which 60 // is used for 'output' which reshapes output from [num_filters, batch, 1] 61 // shape to [num_units, rank, batch] shape. The 2nd shape element is rank. 62 int InferFilterRank(); 63 64 std::vector<std::string> const_node_patterns_; 65 }; 66 67 class SvdfClusterFactory : public ClusterFactoryInterface { 68 public: 69 // Creates a cluster of nodes using a name-based pattern matching approach. It 70 // uses a node as a seed and if its name matches a certain pattern, then it 71 // builds the cluster around that node. 72 // This factory expects nodes which have "SVDF_weights_feature" and 73 // "SVDF_weights_time" pattern in their names (and optionally "SVDF_bias") 74 // and it creates an SVDF Op from them. 75 std::unique_ptr<Cluster> CreateCluster( 76 const tensorflow::NodeDef& node, 77 const tensorflow::GraphDef& graph_def) const; 78 }; 79 80 } // end namespace toco 81 82 #endif // TENSORFLOW_LITE_TOCO_TENSORFLOW_GRAPH_MATCHING_RESOLVE_SVDF_H_ 83