1 /* Copyright 2018 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 
16 #ifndef TENSORFLOW_LITE_DELEGATES_FLEX_TEST_UTIL_H_
17 #define TENSORFLOW_LITE_DELEGATES_FLEX_TEST_UTIL_H_
18 
19 #include "tensorflow/c/c_api_internal.h"
20 #include "tensorflow/lite/kernels/test_util.h"
21 
22 namespace tflite {
23 namespace flex {
24 namespace testing {
25 
26 enum TfOpType {
27   kUnpack,
28   kIdentity,
29   kAdd,
30   kMul,
31   kRfft,
32   kImag,
33   // Represents an op that does not exist in TensorFlow.
34   kNonExistent,
35   // Represents an valid TensorFlow op where the NodeDef is incompatible.
36   kIncompatibleNodeDef,
37 };
38 
39 // This class creates models with TF and TFLite ops. In order to use this class
40 // to test the Flex delegate, implement a function that calls
41 // interpreter->ModifyGraphWithDelegate.
42 class FlexModelTest : public ::testing::Test {
43  public:
FlexModelTest()44   FlexModelTest() {}
~FlexModelTest()45   ~FlexModelTest() {}
46 
47   bool Invoke();
48 
49   // Sets the (typed) tensor's values at the given index.
50   template <typename T>
SetTypedValues(int tensor_index,const std::vector<T> & values)51   void SetTypedValues(int tensor_index, const std::vector<T>& values) {
52     memcpy(interpreter_->typed_tensor<T>(tensor_index), values.data(),
53            values.size() * sizeof(T));
54   }
55 
56   // Returns the (typed) tensor's values at the given index.
57   template <typename T>
GetTypedValues(int tensor_index)58   std::vector<T> GetTypedValues(int tensor_index) {
59     const TfLiteTensor* t = interpreter_->tensor(tensor_index);
60     const T* tdata = interpreter_->typed_tensor<T>(tensor_index);
61     return std::vector<T>(tdata, tdata + t->bytes / sizeof(T));
62   }
63 
64   // Sets the tensor's values at the given index.
SetValues(int tensor_index,const std::vector<float> & values)65   void SetValues(int tensor_index, const std::vector<float>& values) {
66     SetTypedValues<float>(tensor_index, values);
67   }
68   void SetStringValues(int tensor_index, const std::vector<string>& values);
69 
70   // Returns the tensor's values at the given index.
GetValues(int tensor_index)71   std::vector<float> GetValues(int tensor_index) {
72     return GetTypedValues<float>(tensor_index);
73   }
74   std::vector<string> GetStringValues(int tensor_index) const;
75 
76   // Sets the tensor's shape at the given index.
77   void SetShape(int tensor_index, const std::vector<int>& values);
78 
79   // Returns the tensor's shape at the given index.
80   std::vector<int> GetShape(int tensor_index);
81 
82   // Returns the tensor's type at the given index.
83   TfLiteType GetType(int tensor_index);
84 
85   // Returns if the tensor at the given index is dynamic.
86   bool IsDynamicTensor(int tensor_index);
87 
error_reporter()88   const TestErrorReporter& error_reporter() const { return error_reporter_; }
89 
90   // Adds `num_tensor` tensors to the model. `inputs` contains the indices of
91   // the input tensors and `outputs` contains the indices of the output
92   // tensors. All tensors are set to have `type` and `dims`.
93   void AddTensors(int num_tensors, const std::vector<int>& inputs,
94                   const std::vector<int>& outputs, TfLiteType type,
95                   const std::vector<int>& dims);
96 
97   // Set a constant tensor of the given shape, type and buffer at the given
98   // index.
99   void SetConstTensor(int tensor_index, const std::vector<int>& values,
100                       TfLiteType type, const char* buffer, size_t bytes);
101 
102   // Adds a TFLite Mul op. `inputs` contains the indices of the input tensors
103   // and `outputs` contains the indices of the output tensors.
104   void AddTfLiteMulOp(const std::vector<int>& inputs,
105                       const std::vector<int>& outputs);
106 
107   // Adds a TensorFlow op. `inputs` contains the indices of the
108   // input tensors and `outputs` contains the indices of the output tensors.
109   // This function is limited to the set of ops defined in TfOpType.
110   void AddTfOp(TfOpType op, const std::vector<int>& inputs,
111                const std::vector<int>& outputs);
112 
113  protected:
114   std::unique_ptr<Interpreter> interpreter_;
115   TestErrorReporter error_reporter_;
116   std::vector<int> tf_ops_;
117 
118  private:
119   // Helper method to add a TensorFlow op. tflite_names needs to start with
120   // "Flex" in order to work with the Flex delegate.
121   void AddTfOp(const char* tflite_name, const string& tf_name,
122                const string& nodedef_str, const std::vector<int>& inputs,
123                const std::vector<int>& outputs);
124 
125   std::vector<std::vector<uint8_t>> flexbuffers_;
126 
127   int next_op_index_ = 0;
128 };
129 
130 }  // namespace testing
131 }  // namespace flex
132 }  // namespace tflite
133 
134 #endif  // TENSORFLOW_LITE_DELEGATES_FLEX_TEST_UTIL_H_
135