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_PYTHON_OPTIMIZE_CALIBRATION_WRAPPER_H_ 16 #define TENSORFLOW_LITE_PYTHON_OPTIMIZE_CALIBRATION_WRAPPER_H_ 17 18 #include <functional> 19 #include <memory> 20 #include <string> 21 #include <vector> 22 23 // Place `<locale>` before <Python.h> to avoid build failures in macOS. 24 #include <locale> 25 26 // The empty line above is on purpose as otherwise clang-format will 27 // automatically move <Python.h> before <locale>. 28 #include <Python.h> 29 30 #include "tensorflow/lite/interpreter.h" 31 32 // We forward declare TFLite classes here to avoid exposing them to SWIG. 33 namespace tflite { 34 namespace ops { 35 namespace builtin { 36 class BuiltinOpResolver; 37 } // namespace builtin 38 } // namespace ops 39 40 class FlatBufferModel; 41 42 namespace interpreter_wrapper { 43 class PythonErrorReporter; 44 } // namespace interpreter_wrapper 45 46 namespace optimize { 47 namespace calibration { 48 class CalibrationReader; 49 } // namespace calibration 50 } // namespace optimize 51 52 namespace calibration_wrapper { 53 54 PyObject* AddIntermediateTensors(PyObject* data); 55 56 class CalibrationWrapper { 57 public: 58 // SWIG caller takes ownership of pointer. 59 static CalibrationWrapper* CreateWrapperCPPFromBuffer( 60 PyObject* data, const std::vector<std::string>& registerers_by_name, 61 const std::vector<std::function<void(uintptr_t)>>& registerers_by_func, 62 std::string* error_msg); 63 ~CalibrationWrapper(); 64 65 // Allocates the primary subgraph's tensors. 66 PyObject* Prepare(); 67 68 // Allocates the tensors of the given signature, defined by the signature 69 // key. 70 PyObject* Prepare(std::string signature_key); 71 72 // Allocates the primary subgraph's tensors with the given input shapes. 73 PyObject* Prepare(PyObject* input_shapes); 74 75 // Allocates the tensors of the given signature with the given input 76 // shapes, defined by the signature key. 77 PyObject* Prepare(PyObject* input_shapes, std::string signature_key); 78 79 // Sets the given input tensors to the primary subgraph. 80 PyObject* FeedTensor(PyObject* input_value); 81 82 // Sets the given input tensor to the given signature, defined by the 83 // signature key. 84 PyObject* FeedTensor(PyObject* input_value, std::string signature_key); 85 86 // Allows quantizing only the operator that produces the tensor. 87 PyObject* QuantizeModel(int input_py_type, int output_py_type, 88 bool allow_float, int activations_py_type, 89 int bias_py_type); 90 91 // Allows quantizing only the operator that produces the tensor with name 92 // operator_output_name. (This can be used to help debug.). 93 // TODO(suharshs): Allow providing multiple names. 94 PyObject* QuantizeModel(int input_py_type, int output_py_type, 95 bool allow_float, const char* operator_output_name); 96 97 // Disables per-channel quantization, can be used to produce smaller 98 // models but may cause accuracy issues. 99 PyObject* QuantizeModel(int input_py_type, int output_py_type, 100 bool allow_float, int activations_py_type, 101 int bias_py_type, bool disable_per_channel); 102 103 // Writes the in-memory calibration results to the model flatbuffer. The 104 // produced model is as same as the original input model, but the min/max 105 // in the quantization field. 106 PyObject* Calibrate(); 107 108 private: 109 // CalibrationWrapper is not copyable or assignable. We avoid the use of 110 // CalibrationWrapper() = delete here for SWIG compatibility. 111 CalibrationWrapper( 112 std::unique_ptr<tflite::Interpreter> interpreter, 113 std::unique_ptr<tflite::ops::builtin::BuiltinOpResolver> resolver, 114 std::unique_ptr<tflite::interpreter_wrapper::PythonErrorReporter> 115 error_reporter, 116 std::unique_ptr<tflite::FlatBufferModel> model, 117 std::unique_ptr<tflite::optimize::calibration::CalibrationReader> reader, 118 std::unique_ptr<std::string> model_str_); 119 120 CalibrationWrapper(const CalibrationWrapper& rhs); 121 122 PyObject* SetTensor(int index, PyObject* value); 123 PyObject* SetTensor(int index, PyObject* value, std::string signature_key); 124 125 std::unique_ptr<tflite::Interpreter> interpreter_; 126 std::unique_ptr<tflite::interpreter_wrapper::PythonErrorReporter> 127 error_reporter_; 128 std::unique_ptr<tflite::ops::builtin::BuiltinOpResolver> resolver_; 129 std::unique_ptr<tflite::FlatBufferModel> model_; 130 std::unique_ptr<tflite::optimize::calibration::CalibrationReader> reader_; 131 std::unique_ptr<std::string> model_str_; 132 }; 133 134 } // namespace calibration_wrapper 135 } // namespace tflite 136 137 #endif // TENSORFLOW_LITE_PYTHON_OPTIMIZE_CALIBRATION_WRAPPER_H_ 138