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_TFLITE_OPERATOR_H_ 16 #define TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_OPERATOR_H_ 17 18 #include "flatbuffers/flatbuffers.h" 19 #include "tensorflow/contrib/lite/schema/schema_generated.h" 20 #include "tensorflow/contrib/lite/toco/model.h" 21 22 namespace toco { 23 24 namespace tflite { 25 26 class BaseOperator; 27 28 // Return a map contained all knwo TF Lite Operators, keyed by their names. 29 std::map<string, std::unique_ptr<BaseOperator>> BuildOperatorByNameMap(); 30 31 // Return a map contained all knwo TF Lite Operators, keyed by the type of 32 // their tf.mini counterparts. 33 std::map<OperatorType, std::unique_ptr<BaseOperator>> BuildOperatorByTypeMap(); 34 35 // These are the flatbuffer types for custom and builtin options. 36 using CustomOptions = flatbuffers::Vector<uint8_t>; 37 using BuiltinOptions = void; 38 39 // A simple wrapper around the flatbuffer objects used to describe options that 40 // configure operators. 41 struct Options { 42 // Build custom options. CustomOptions43 static Options Custom(flatbuffers::Offset<CustomOptions> offset) { 44 return {::tflite::BuiltinOptions_NONE, 0, offset}; 45 } 46 47 // Build builtin options of the given type. BuiltinOptions48 static Options Builtin(::tflite::BuiltinOptions type, 49 flatbuffers::Offset<BuiltinOptions> offset) { 50 return {type, offset, 0}; 51 } 52 53 ::tflite::BuiltinOptions type; 54 flatbuffers::Offset<BuiltinOptions> builtin; 55 flatbuffers::Offset<CustomOptions> custom; 56 }; 57 58 // A BaseOperator encapsulates the relationship between operators in tf.mini 59 // and TF lite, and provides methods for converting between those two formats. 60 class BaseOperator { 61 public: 62 // Build an operator with the given TF Lite name and tf.mini type. BaseOperator(const string & name,OperatorType type)63 BaseOperator(const string& name, OperatorType type) 64 : name_(name), type_(type) {} 65 virtual ~BaseOperator() = default; 66 name()67 string name() const { return name_; } type()68 OperatorType type() const { return type_; } 69 70 // Given a tf.mini operator, create the corresponding flatbuffer options and 71 // return their offsets. 72 virtual Options Serialize(const Operator& op, 73 flatbuffers::FlatBufferBuilder* builder) const = 0; 74 75 // Read TF Lite options and create the appropriate tf.mini operator. 76 virtual std::unique_ptr<Operator> Deserialize( 77 const BuiltinOptions* builtin_options, 78 const CustomOptions* custom_options) const = 0; 79 80 private: 81 string name_; 82 OperatorType type_; 83 }; 84 85 } // namespace tflite 86 87 } // namespace toco 88 89 #endif // TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_OPERATOR_H_ 90