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