compiler/tflchef/tflite/src/TFliteImport.h

   1 /*
   2  * Copyright (c) 2018 Samsung Electronics Co., Ltd. All Rights Reserved
   3  *
   4  * Licensed under the Apache License, Version 2.0 (the "License");
   5  * you may not use this file except in compliance with the License.
   6  * You may obtain a copy of the License at
   7  *
   8  *    http://www.apache.org/licenses/LICENSE-2.0
   9  *
  10  * Unless required by applicable law or agreed to in writing, software
  11  * distributed under the License is distributed on an "AS IS" BASIS,
  12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  * See the License for the specific language governing permissions and
  14  * limitations under the License.
  15  */
  16
  17 #ifndef __TFLITE_IMPORT_H__
  18 #define __TFLITE_IMPORT_H__
  19
  20 #include <mio/tflite/schema_generated.h>
  21
  22 #include <tflchef.pb.h>
  23
  24 #include <map>
  25 #include <vector>
  26
  27 namespace tflchef
  28 {
  29
  30 using TFliteSubGraphs_t = flatbuffers::Vector<flatbuffers::Offset<tflite::SubGraph>>;
  31 using TFliteTensors_t = flatbuffers::Vector<flatbuffers::Offset<tflite::Tensor>>;
  32 using TFliteBuffers_t = flatbuffers::Vector<flatbuffers::Offset<tflite::Buffer>>;
  33 using TFliteOperators_t = flatbuffers::Vector<flatbuffers::Offset<tflite::Operator>>;
  34
  35 const char *tensor_type(const tflite::Tensor *tensor);
  36 const char *tensor_name(const tflite::Tensor *tensor);
  37 bool is_valid(const tflite::OperatorCode *opcode);
  38 bool is_custom(const tflite::OperatorCode *opcode);
  39
  40 /**
  41  * @brief Loads TF lite file and provides helpers to access attributes
  42  */
  43 class TFliteImport
  44 {
  45 public:
  46   TFliteImport(const tflite::Model *model);
  47
  48   TFliteImport() = delete;
  49
  50 public:
  51   bool select_sub_graph(uint32_t subgraph);
  52
  53 public:
  54   const TFliteBuffers_t *buffers() { return _buffers; }
  55   const TFliteTensors_t *tensors() { return _tensors; }
  56   const TFliteOperators_t *operators() { return _operators; }
  57   const std::vector<int32_t> &inputs() const { return _inputs; }
  58   const std::vector<int32_t> &outputs() const { return _outputs; }
  59
  60   uint32_t num_subgraph() const { return _subgraphs->Length(); }
  61
  62   tflite::BuiltinOperator builtin_code(const tflite::Operator *op) const;
  63   std::string opcode_name(const tflite::Operator *op) const;
  64   size_t buffer_info(const tflite::Tensor *tensor, const uint8_t **buff_data);
  65
  66   /**
  67    * @brief This will record the tensor by index, if it needs filler option,
  68    *        such as kernel, bias.
  69    */
  70   void set_tensor_filler(uint32_t tensor_index) { _tensor_filler[tensor_index] = true; }
  71
  72   /**
  73    * @brief This will store int32 filler values such as reshape information for the tensor
  74    */
  75   void set_tensor_filler(uint32_t tensor_index, std::vector<int32_t> &expvalues)
  76   {
  77     _tensor_filler_vint32[tensor_index] = expvalues;
  78   }
  79
  80   void set_tensor_filler(uint32_t tensor_index, std::vector<float> &expvalues)
  81   {
  82     _tensor_filler_vfloat[tensor_index] = expvalues;
  83   }
  84
  85   /**
  86    * @brief This will return true if the tensor by index, needs a filler option.
  87    */
  88   bool get_tensor_filler(uint32_t tensor_index)
  89   {
  90     auto it = _tensor_filler.find(tensor_index);
  91     if (it != _tensor_filler.end())
  92     {
  93       return it->second;
  94     }
  95     return false;
  96   }
  97
  98   /**
  99    * @brief This will return true if the tensor by index, needs a int array filler option.
 100    */
 101   bool get_tensor_filler(uint32_t tensor_index, std::vector<int32_t> &expvalues)
 102   {
 103     auto it = _tensor_filler_vint32.find(tensor_index);
 104     if (it != _tensor_filler_vint32.end())
 105     {
 106       expvalues = it->second;
 107       return true;
 108     }
 109     return false;
 110   }
 111
 112   bool get_tensor_filler(uint32_t tensor_index, std::vector<float> &expvalues)
 113   {
 114     auto it = _tensor_filler_vfloat.find(tensor_index);
 115     if (it != _tensor_filler_vfloat.end())
 116     {
 117       expvalues = it->second;
 118       return true;
 119     }
 120     return false;
 121   }
 122
 123 private:
 124   const TFliteSubGraphs_t *_subgraphs{nullptr};
 125   const TFliteBuffers_t *_buffers{nullptr};
 126   const TFliteTensors_t *_tensors{nullptr};
 127   const TFliteOperators_t *_operators{nullptr};
 128
 129   std::vector<const tflite::OperatorCode *> _op_codes{};
 130   std::vector<int32_t> _inputs{};
 131   std::vector<int32_t> _outputs{};
 132
 133   std::map<uint32_t, bool> _tensor_filler{};
 134   std::map<uint32_t, std::vector<int32_t>> _tensor_filler_vint32{};
 135   std::map<uint32_t, std::vector<float>> _tensor_filler_vfloat{};
 136 };
 137
 138 } // namespace tflchef
 139
 140 #endif // __TFLITE_IMPORT_H__