1 /* Copyright 2019 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_EXPERIMENTAL_DELEGATES_COREML_BUILDERS_OP_BUILDER_H_ 16 #define TENSORFLOW_LITE_EXPERIMENTAL_DELEGATES_COREML_BUILDERS_OP_BUILDER_H_ 17 18 #include <functional> 19 #include <string> 20 21 #include "mlmodel/format/Model.pb.h" 22 #include "mlmodel/format/NeuralNetwork.pb.h" 23 #include "tensorflow/lite/c/common.h" 24 25 namespace tflite { 26 namespace delegates { 27 namespace coreml { 28 class OpBuilder; 29 30 // A class represents an ID in the coreML graph. 31 // A node is represented by a pair (node_id, and output_index) 32 // API is experimental and subject to change. 33 class TensorID { 34 public: TensorID()35 TensorID() {} TensorID(int node,int output_id)36 TensorID(int node, int output_id) : node_(node), output_id_(output_id) {} 37 38 std::string ToString() const; 39 40 int NodeID() const; 41 42 int OutputID() const; 43 44 private: 45 int node_ = -1; 46 int output_id_ = -1; 47 }; 48 49 // Builder for the whole graph. 50 // All op builders should be added using AddBuilder 51 // and then BuildModel should be called to return the CoreML generated. 52 // 53 // API is experimental and subject to change. 54 class GraphBuilder { 55 public: GraphBuilder(int coreml_version)56 explicit GraphBuilder(int coreml_version) : coreml_version_(coreml_version) {} 57 58 // Returns pointer to the created builder. Ownership still belongs 59 // to the GraphBuilder. 60 OpBuilder* AddBuilder(int builtin_code, const TfLiteNode* node); 61 62 // Returns pointer to the created builder with op builder function provided. 63 OpBuilder* AddBuilder(const std::function<OpBuilder*(GraphBuilder*)>& builder, 64 const TfLiteNode* node); 65 66 // Builds Model instance and returns it. 67 CoreML::Specification::Model* BuildModel(); 68 69 // Returns string representing tensor 'tensor_id' in coreML. 70 // tensor_id should have been added before calling this method. 71 std::string GetTensorName(int tensor_id); 72 73 // Returns Core ML Tensor ID for TFL 'tensor_id'. 74 // tensor_id should have been added before calling this method. 75 const TensorID GetTensorID(int tensor_id); 76 77 void AddTensorWithID(int tf_tensor_id, const TensorID& tensor_id); 78 79 // Return true if this tensor was added before to the graph. 80 bool HasTensor(int tflite_tensor_index); 81 // Return if this tensor is used in the graph (not as data). 82 // This information is used to mark constant tensors that are used as input. 83 bool IsTensorUsed(int tflite_tensor_index); 84 85 const int coreml_version_; 86 87 private: 88 std::vector<std::unique_ptr<OpBuilder>> builders_; 89 // Index in the vector is the tflite_tensor_index, the value 90 // is the ID in the coreml graph. 91 std::vector<TensorID> tensors_; 92 std::vector<bool> used_tensor_; 93 }; 94 95 // Interface for all op layers 96 // API is experimental and subject to change. 97 class OpBuilder { 98 public: OpBuilder(GraphBuilder * graph_builder)99 explicit OpBuilder(GraphBuilder* graph_builder) 100 : graph_builder_(graph_builder) {} ~OpBuilder()101 virtual ~OpBuilder() {} 102 103 // Returns the Layer this builder responsible for. 104 // Ownership is transferred to caller. 105 virtual CoreML::Specification::NeuralNetworkLayer* Build(); 106 107 // Associates TfLite input tensors to Core ML layer's inputs and properties. 108 // Verification for input constraints should happen here. 109 virtual TfLiteStatus RegisterInputs(const TfLiteIntArray* inputs, 110 TfLiteContext* context) = 0; 111 112 // Associates TFLite output tensor with the node's output. If the OpBuilder 113 // has subgraphs, The final output of that subgraph should be associated with 114 // the output tensor. 115 virtual TfLiteStatus RegisterOutputs(const TfLiteIntArray* outputs, 116 TfLiteContext* context) = 0; 117 118 // Adds additional required OpBuilders, and populate builder_output_ with 119 // Actual output that corresponds to output tensor of TFL Node. 120 // Clients need to override this in cases where the nodes can be used for 121 // composing other ops. For example, Relu6 in TfLite can be converted to 122 // Relu -> Threshold -> Neg. 123 // TODO(b/147211734): have this called automatically when necessary. 124 virtual TfLiteStatus PopulateSubgraph(TfLiteContext* context); 125 126 virtual const std::string& DebugName() = 0; 127 128 void SetBuiltinData(void* builtin_data); 129 130 void SetNodeID(int id); 131 132 void SetTfLiteNode(const TfLiteNode* node); 133 134 int GetID() const; 135 136 // Adds input with tensor name. 137 void AddInput(const std::string& input_name); 138 139 // Adds input with CoreML tensor ID. 140 void AddInput(const TensorID& input_id); 141 142 // Adds input with TF Lite tensor ID. 143 // TODO(taeheej): cleanup AddInput use cases and used tensor tracking. 144 void AddInput(int tf_input_id); 145 146 // Simply adds new output to the underlying layer. 147 TensorID AddOutput(); 148 149 // Should set builder_output_ (if unset) and return it as the output of 150 // this node. To be used by clients that needs the output of the node. 151 virtual TensorID GetOutput(TfLiteContext* context); 152 153 protected: 154 // Sets layer's name. 155 void SetDebugName(const char* layer_name, int id); 156 157 GraphBuilder* graph_builder_ = nullptr; 158 // Data needed by this node. 159 void* builtin_data_ = nullptr; 160 int node_id_ = -1; 161 int num_outputs_ = 0; 162 const TfLiteNode* tflite_node_ = nullptr; 163 TensorID builder_output_; 164 std::string debug_name_; 165 std::unique_ptr<CoreML::Specification::NeuralNetworkLayer> layer_; 166 }; 167 168 } // namespace coreml 169 } // namespace delegates 170 } // namespace tflite 171 172 #endif // TENSORFLOW_LITE_EXPERIMENTAL_DELEGATES_COREML_BUILDERS_OP_BUILDER_H_ 173