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_TOOLS_EVALUATION_STAGES_TFLITE_INFERENCE_STAGE_H_ 16 #define TENSORFLOW_LITE_TOOLS_EVALUATION_STAGES_TFLITE_INFERENCE_STAGE_H_ 17 18 #include <stdint.h> 19 20 #include <vector> 21 22 #include "tensorflow/core/util/stats_calculator.h" 23 #include "tensorflow/lite/c/common.h" 24 #include "tensorflow/lite/interpreter.h" 25 #include "tensorflow/lite/kernels/register.h" 26 #include "tensorflow/lite/model.h" 27 #include "tensorflow/lite/tools/evaluation/evaluation_delegate_provider.h" 28 #include "tensorflow/lite/tools/evaluation/evaluation_stage.h" 29 #include "tensorflow/lite/tools/evaluation/proto/evaluation_config.pb.h" 30 31 namespace tflite { 32 namespace evaluation { 33 34 struct TfLiteModelInfo { 35 std::vector<const TfLiteTensor*> inputs; 36 std::vector<const TfLiteTensor*> outputs; 37 }; 38 39 // EvaluationStage to run inference using TFLite. 40 class TfliteInferenceStage : public EvaluationStage { 41 public: TfliteInferenceStage(const EvaluationStageConfig & config)42 explicit TfliteInferenceStage(const EvaluationStageConfig& config) 43 : EvaluationStage(config) {} 44 Init()45 TfLiteStatus Init() override { return Init(nullptr); } 46 TfLiteStatus Init(const DelegateProviders* delegate_providers); 47 48 TfLiteStatus Run() override; 49 50 // EvaluationStageMetrics.num_runs denotes the number of inferences run. 51 EvaluationStageMetrics LatestMetrics() override; 52 ~TfliteInferenceStage()53 ~TfliteInferenceStage() override {} 54 55 // Call before Run(). 56 // This class does not take ownership of raw_input_ptrs. SetInputs(const std::vector<void * > & raw_input_ptrs)57 void SetInputs(const std::vector<void*>& raw_input_ptrs) { 58 inputs_ = &raw_input_ptrs; 59 } 60 61 // Resize input tensors with given shapes. 62 TfLiteStatus ResizeInputs(const std::vector<std::vector<int>>& shapes); 63 64 // Applies provided delegate to the underlying TFLite Interpreter. 65 TfLiteStatus ApplyCustomDelegate(Interpreter::TfLiteDelegatePtr delegate); 66 67 // Read-only view of a TfliteModelInfo. TfliteInferenceStage retains 68 // ownership. 69 // Only available after Init is done. GetModelInfo()70 const TfLiteModelInfo* GetModelInfo() const { return &model_info_; } 71 72 // Provides a read-only view to the model's output tensor(s). Retains 73 // ownership of object. GetOutputs()74 const std::vector<void*>* GetOutputs() const { return &outputs_; } 75 76 private: 77 // Sets model_info_ & outputs_ after interpreter tensors are (re)allocated. 78 void UpdateModelInfo(); 79 80 std::unique_ptr<FlatBufferModel> model_; 81 std::unique_ptr<ops::builtin::BuiltinOpResolver> resolver_; 82 std::unique_ptr<Interpreter> interpreter_; 83 std::vector<Interpreter::TfLiteDelegatePtr> delegates_; 84 85 TfLiteModelInfo model_info_; 86 const std::vector<void*>* inputs_ = nullptr; 87 std::vector<void*> outputs_; 88 89 tensorflow::Stat<int64_t> latency_stats_; 90 }; 91 92 } // namespace evaluation 93 } // namespace tflite 94 95 #endif // TENSORFLOW_LITE_TOOLS_EVALUATION_STAGES_TFLITE_INFERENCE_STAGE_H_ 96