| /* |
| * Copyright (C) 2018 The Android Open Source Project |
| * |
| * Licensed under the Apache License, Version 2.0 (the "License"); |
| * you may not use this file except in compliance with the License. |
| * You may obtain a copy of the License at |
| * |
| * http://www.apache.org/licenses/LICENSE-2.0 |
| * |
| * Unless required by applicable law or agreed to in writing, software |
| * distributed under the License is distributed on an "AS IS" BASIS, |
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| * See the License for the specific language governing permissions and |
| * limitations under the License. |
| */ |
| |
| #include "utils/tflite-model-executor.h" |
| |
| #include "utils/base/logging.h" |
| #include "tensorflow/lite/kernels/register.h" |
| #include "tensorflow/lite/schema/schema_generated.h" |
| |
| // Forward declaration of custom TensorFlow Lite ops for registration. |
| namespace tflite { |
| namespace ops { |
| namespace builtin { |
| TfLiteRegistration* Register_GELU(); |
| TfLiteRegistration* Register_ADD(); |
| TfLiteRegistration* Register_CONCATENATION(); |
| TfLiteRegistration* Register_CONV_2D(); |
| TfLiteRegistration* Register_DEPTHWISE_CONV_2D(); |
| TfLiteRegistration* Register_AVERAGE_POOL_2D(); |
| TfLiteRegistration* Register_EQUAL(); |
| TfLiteRegistration* Register_FULLY_CONNECTED(); |
| TfLiteRegistration* Register_GREATER_EQUAL(); |
| TfLiteRegistration* Register_L2_NORMALIZATION(); |
| TfLiteRegistration* Register_MUL(); |
| TfLiteRegistration* Register_RESHAPE(); |
| TfLiteRegistration* Register_REDUCE_MAX(); |
| TfLiteRegistration* Register_REDUCE_MIN(); |
| TfLiteRegistration* Register_REDUCE_ANY(); |
| TfLiteRegistration* Register_SOFTMAX(); |
| TfLiteRegistration* Register_GATHER(); |
| TfLiteRegistration* Register_GATHER_ND(); |
| TfLiteRegistration* Register_IF(); |
| TfLiteRegistration* Register_ROUND(); |
| TfLiteRegistration* Register_ZEROS_LIKE(); |
| TfLiteRegistration* Register_TRANSPOSE(); |
| TfLiteRegistration* Register_SUB(); |
| TfLiteRegistration* Register_DIV(); |
| TfLiteRegistration* Register_STRIDED_SLICE(); |
| TfLiteRegistration* Register_EXP(); |
| TfLiteRegistration* Register_TOPK_V2(); |
| TfLiteRegistration* Register_SLICE(); |
| TfLiteRegistration* Register_SPLIT(); |
| TfLiteRegistration* Register_CAST(); |
| TfLiteRegistration* Register_MAXIMUM(); |
| TfLiteRegistration* Register_MINIMUM(); |
| TfLiteRegistration* Register_NEG(); |
| TfLiteRegistration* Register_SLICE(); |
| TfLiteRegistration* Register_LOG(); |
| TfLiteRegistration* Register_LOGISTIC(); |
| TfLiteRegistration* Register_SUM(); |
| TfLiteRegistration* Register_PACK(); |
| TfLiteRegistration* Register_DEQUANTIZE(); |
| TfLiteRegistration* Register_MEAN(); |
| TfLiteRegistration* Register_LESS(); |
| TfLiteRegistration* Register_TILE(); |
| TfLiteRegistration* Register_SQUARED_DIFFERENCE(); |
| TfLiteRegistration* Register_RSQRT(); |
| TfLiteRegistration* Register_LOG_SOFTMAX(); |
| TfLiteRegistration* Register_WHERE(); |
| TfLiteRegistration* Register_ONE_HOT(); |
| TfLiteRegistration* Register_POW(); |
| TfLiteRegistration* Register_TANH(); |
| TfLiteRegistration* Register_UNIQUE(); |
| TfLiteRegistration* Register_REDUCE_PROD(); |
| TfLiteRegistration* Register_SHAPE(); |
| TfLiteRegistration* Register_NOT_EQUAL(); |
| TfLiteRegistration* Register_CUMSUM(); |
| TfLiteRegistration* Register_EXPAND_DIMS(); |
| TfLiteRegistration* Register_FILL(); |
| TfLiteRegistration* Register_PADV2(); |
| TfLiteRegistration* Register_EMBEDDING_LOOKUP(); |
| TfLiteRegistration* Register_GREATER(); |
| } // namespace builtin |
| } // namespace ops |
| } // namespace tflite |
| |
| #ifdef TC3_WITH_ACTIONS_OPS |
| #include "utils/tflite/blacklist.h" |
| #include "utils/tflite/dist_diversification.h" |
| #include "utils/tflite/string_projection.h" |
| #include "utils/tflite/text_encoder.h" |
| #include "utils/tflite/text_encoder3s.h" |
| #include "utils/tflite/token_encoder.h" |
| |
| namespace tflite { |
| namespace ops { |
| namespace custom { |
| TfLiteRegistration* Register_SENTENCEPIECE_TOKENIZER(); |
| TfLiteRegistration* Register_RAGGED_TENSOR_TO_TENSOR(); |
| TfLiteRegistration* Register_RAGGED_RANGE(); |
| TfLiteRegistration* Register_RANDOM_UNIFORM(); |
| } // namespace custom |
| } // namespace ops |
| } // namespace tflite |
| |
| void RegisterSelectedOps(tflite::MutableOpResolver* resolver) { |
| resolver->AddBuiltin(tflite::BuiltinOperator_ADD, |
| tflite::ops::builtin::Register_ADD(), |
| /*min_version=*/1, |
| /*max_version=*/2); |
| resolver->AddBuiltin(tflite::BuiltinOperator_CONCATENATION, |
| tflite::ops::builtin::Register_CONCATENATION(), |
| /*min_version=*/1, |
| /*max_version=*/2); |
| resolver->AddBuiltin(tflite::BuiltinOperator_CONV_2D, |
| tflite::ops::builtin::Register_CONV_2D(), |
| /*min_version=*/1, |
| /*max_version=*/5); |
| resolver->AddBuiltin(tflite::BuiltinOperator_DEPTHWISE_CONV_2D, |
| tflite::ops::builtin::Register_DEPTHWISE_CONV_2D(), |
| /*min_version=*/1, |
| /*max_version=*/6); |
| resolver->AddBuiltin(tflite::BuiltinOperator_AVERAGE_POOL_2D, |
| tflite::ops::builtin::Register_AVERAGE_POOL_2D(), |
| /*min_version=*/1, |
| /*max_version=*/1); |
| resolver->AddBuiltin(::tflite::BuiltinOperator_EQUAL, |
| ::tflite::ops::builtin::Register_EQUAL()); |
| |
| resolver->AddBuiltin(tflite::BuiltinOperator_FULLY_CONNECTED, |
| tflite::ops::builtin::Register_FULLY_CONNECTED(), |
| /*min_version=*/1, |
| /*max_version=*/9); |
| resolver->AddBuiltin(::tflite::BuiltinOperator_GREATER_EQUAL, |
| ::tflite::ops::builtin::Register_GREATER_EQUAL()); |
| resolver->AddBuiltin(tflite::BuiltinOperator_L2_NORMALIZATION, |
| tflite::ops::builtin::Register_L2_NORMALIZATION(), |
| /*min_version=*/1, |
| /*max_version=*/2); |
| resolver->AddBuiltin(tflite::BuiltinOperator_MUL, |
| tflite::ops::builtin::Register_MUL()); |
| resolver->AddBuiltin(tflite::BuiltinOperator_RESHAPE, |
| tflite::ops::builtin::Register_RESHAPE()); |
| resolver->AddBuiltin(::tflite::BuiltinOperator_REDUCE_MAX, |
| ::tflite::ops::builtin::Register_REDUCE_MAX()); |
| resolver->AddBuiltin(::tflite::BuiltinOperator_REDUCE_MIN, |
| ::tflite::ops::builtin::Register_REDUCE_MIN()); |
| resolver->AddBuiltin(::tflite::BuiltinOperator_REDUCE_ANY, |
| ::tflite::ops::builtin::Register_REDUCE_ANY()); |
| resolver->AddBuiltin(tflite::BuiltinOperator_SOFTMAX, |
| tflite::ops::builtin::Register_SOFTMAX(), |
| /*min_version=*/1, |
| /*max_version=*/2); |
| resolver->AddBuiltin(tflite::BuiltinOperator_GATHER, |
| tflite::ops::builtin::Register_GATHER(), |
| /*min_version=*/1, |
| /*max_version=*/2); |
| resolver->AddBuiltin(::tflite::BuiltinOperator_GATHER_ND, |
| ::tflite::ops::builtin::Register_GATHER_ND(), |
| /*version=*/2); |
| resolver->AddBuiltin(::tflite::BuiltinOperator_IF, |
| ::tflite::ops::builtin::Register_IF()), |
| resolver->AddBuiltin(::tflite::BuiltinOperator_ROUND, |
| ::tflite::ops::builtin::Register_ROUND()); |
| resolver->AddBuiltin(::tflite::BuiltinOperator_ZEROS_LIKE, |
| ::tflite::ops::builtin::Register_ZEROS_LIKE()); |
| resolver->AddBuiltin(tflite::BuiltinOperator_TRANSPOSE, |
| tflite::ops::builtin::Register_TRANSPOSE(), |
| /*min_version=*/1, |
| /*max_version=*/2); |
| resolver->AddBuiltin(tflite::BuiltinOperator_SUB, |
| tflite::ops::builtin::Register_SUB(), |
| /*min_version=*/1, |
| /*max_version=*/2); |
| resolver->AddBuiltin(tflite::BuiltinOperator_DIV, |
| tflite::ops::builtin::Register_DIV()); |
| resolver->AddBuiltin(tflite::BuiltinOperator_STRIDED_SLICE, |
| tflite::ops::builtin::Register_STRIDED_SLICE(), |
| /*min_version=*/1, |
| /*max_version=*/2); |
| resolver->AddBuiltin(tflite::BuiltinOperator_EXP, |
| tflite::ops::builtin::Register_EXP()); |
| resolver->AddBuiltin(tflite::BuiltinOperator_TOPK_V2, |
| tflite::ops::builtin::Register_TOPK_V2(), |
| /*min_version=*/1, |
| /*max_version=*/2); |
| resolver->AddBuiltin(tflite::BuiltinOperator_SLICE, |
| tflite::ops::builtin::Register_SLICE(), |
| /*min_version=*/1, |
| /*max_version=*/3); |
| resolver->AddBuiltin(tflite::BuiltinOperator_SPLIT, |
| tflite::ops::builtin::Register_SPLIT(), |
| /*min_version=*/1, |
| /*max_version=*/3); |
| resolver->AddBuiltin(tflite::BuiltinOperator_CAST, |
| tflite::ops::builtin::Register_CAST()); |
| resolver->AddBuiltin(tflite::BuiltinOperator_MAXIMUM, |
| tflite::ops::builtin::Register_MAXIMUM(), |
| /*min_version=*/1, |
| /*max_version=*/2); |
| resolver->AddBuiltin(tflite::BuiltinOperator_MINIMUM, |
| tflite::ops::builtin::Register_MINIMUM(), |
| /*min_version=*/1, |
| /*max_version=*/2); |
| resolver->AddBuiltin(tflite::BuiltinOperator_NEG, |
| tflite::ops::builtin::Register_NEG()); |
| resolver->AddBuiltin(tflite::BuiltinOperator_SLICE, |
| tflite::ops::builtin::Register_SLICE(), |
| /*min_version=*/1, |
| /*max_version=*/2); |
| resolver->AddBuiltin(tflite::BuiltinOperator_LOG, |
| tflite::ops::builtin::Register_LOG()); |
| resolver->AddBuiltin(tflite::BuiltinOperator_LOGISTIC, |
| tflite::ops::builtin::Register_LOGISTIC()); |
| resolver->AddBuiltin(tflite::BuiltinOperator_SUM, |
| tflite::ops::builtin::Register_SUM()); |
| resolver->AddBuiltin(tflite::BuiltinOperator_PACK, |
| tflite::ops::builtin::Register_PACK(), |
| /*min_version=*/1, |
| /*max_version=*/2); |
| resolver->AddBuiltin(tflite::BuiltinOperator_DEQUANTIZE, |
| tflite::ops::builtin::Register_DEQUANTIZE(), |
| /*min_version=*/1, |
| /*max_version=*/2); |
| resolver->AddBuiltin(tflite::BuiltinOperator_MEAN, |
| tflite::ops::builtin::Register_MEAN()); |
| resolver->AddBuiltin(tflite::BuiltinOperator_LESS, |
| tflite::ops::builtin::Register_LESS()); |
| resolver->AddBuiltin(tflite::BuiltinOperator_TILE, |
| tflite::ops::builtin::Register_TILE()); |
| resolver->AddBuiltin(tflite::BuiltinOperator_SQUARED_DIFFERENCE, |
| tflite::ops::builtin::Register_SQUARED_DIFFERENCE()); |
| resolver->AddBuiltin(tflite::BuiltinOperator_RSQRT, |
| tflite::ops::builtin::Register_RSQRT()); |
| resolver->AddBuiltin(tflite::BuiltinOperator_LOG_SOFTMAX, |
| tflite::ops::builtin::Register_LOG_SOFTMAX()); |
| resolver->AddBuiltin(::tflite::BuiltinOperator_WHERE, |
| ::tflite::ops::builtin::Register_WHERE()); |
| resolver->AddBuiltin(tflite::BuiltinOperator_ONE_HOT, |
| tflite::ops::builtin::Register_ONE_HOT(), |
| /*min_version=*/1, |
| /*max_version=*/1); |
| resolver->AddBuiltin(tflite::BuiltinOperator_POW, |
| tflite::ops::builtin::Register_POW(), |
| /*min_version=*/1, |
| /*max_version=*/1); |
| resolver->AddBuiltin(tflite::BuiltinOperator_TANH, |
| tflite::ops::builtin::Register_TANH(), |
| /*min_version=*/1, |
| /*max_version=*/1); |
| resolver->AddBuiltin(::tflite::BuiltinOperator_UNIQUE, |
| ::tflite::ops::builtin::Register_UNIQUE()); |
| resolver->AddBuiltin(::tflite::BuiltinOperator_REDUCE_PROD, |
| ::tflite::ops::builtin::Register_REDUCE_PROD()); |
| resolver->AddBuiltin(::tflite::BuiltinOperator_SHAPE, |
| ::tflite::ops::builtin::Register_SHAPE()); |
| resolver->AddBuiltin(::tflite::BuiltinOperator_NOT_EQUAL, |
| ::tflite::ops::builtin::Register_NOT_EQUAL()); |
| resolver->AddBuiltin(::tflite::BuiltinOperator_CUMSUM, |
| ::tflite::ops::builtin::Register_CUMSUM()); |
| resolver->AddBuiltin(::tflite::BuiltinOperator_EXPAND_DIMS, |
| ::tflite::ops::builtin::Register_EXPAND_DIMS()); |
| resolver->AddBuiltin(::tflite::BuiltinOperator_FILL, |
| ::tflite::ops::builtin::Register_FILL()); |
| resolver->AddBuiltin(::tflite::BuiltinOperator_PADV2, |
| ::tflite::ops::builtin::Register_PADV2()); |
| resolver->AddBuiltin(::tflite::BuiltinOperator_EMBEDDING_LOOKUP, |
| ::tflite::ops::builtin::Register_EMBEDDING_LOOKUP(), |
| /* min_version=*/1, |
| /*max_version=*/3); |
| resolver->AddBuiltin(::tflite::BuiltinOperator_GREATER, |
| ::tflite::ops::builtin::Register_GREATER()); |
| resolver->AddBuiltin(::tflite::BuiltinOperator_GELU, |
| ::tflite::ops::builtin::Register_GELU()); |
| } |
| #else |
| void RegisterSelectedOps(tflite::MutableOpResolver* resolver) { |
| resolver->AddBuiltin(tflite::BuiltinOperator_FULLY_CONNECTED, |
| tflite::ops::builtin::Register_FULLY_CONNECTED()); |
| } |
| #endif // TC3_WITH_ACTIONS_OPS |
| |
| namespace libtextclassifier3 { |
| |
| std::unique_ptr<tflite::OpResolver> BuildOpResolver() { |
| return BuildOpResolver([](tflite::MutableOpResolver* mutable_resolver) {}); |
| } |
| |
| std::unique_ptr<tflite::OpResolver> BuildOpResolver( |
| const std::function<void(tflite::MutableOpResolver*)>& customize_fn) { |
| #ifdef TC3_USE_SELECTIVE_REGISTRATION |
| std::unique_ptr<tflite::MutableOpResolver> resolver( |
| new tflite::MutableOpResolver); |
| RegisterSelectedOps(resolver.get()); |
| #else |
| std::unique_ptr<tflite::ops::builtin::BuiltinOpResolver> resolver( |
| new tflite::ops::builtin::BuiltinOpResolver); |
| #endif |
| #ifdef TC3_WITH_ACTIONS_OPS |
| resolver->AddCustom("DistanceDiversification", |
| tflite::ops::custom::Register_DISTANCE_DIVERSIFICATION()); |
| resolver->AddCustom("TextEncoder", |
| tflite::ops::custom::Register_TEXT_ENCODER()); |
| resolver->AddCustom("TextEncoder3S", |
| tflite::ops::custom::Register_TEXT_ENCODER3S()); |
| resolver->AddCustom("TokenEncoder", |
| tflite::ops::custom::Register_TOKEN_ENCODER()); |
| resolver->AddCustom( |
| "TFSentencepieceTokenizeOp", |
| ::tflite::ops::custom::Register_SENTENCEPIECE_TOKENIZER()); |
| resolver->AddCustom("RaggedRange", |
| ::tflite::ops::custom::Register_RAGGED_RANGE()); |
| resolver->AddCustom( |
| "RaggedTensorToTensor", |
| ::tflite::ops::custom::Register_RAGGED_TENSOR_TO_TENSOR()); |
| resolver->AddCustom( |
| "STRING_PROJECTION", |
| ::tflite::ops::custom::libtextclassifier3::Register_STRING_PROJECTION()); |
| resolver->AddCustom( |
| "BLACKLIST", |
| ::tflite::ops::custom::libtextclassifier3::Register_BLACKLIST()); |
| resolver->AddCustom("RandomUniform", |
| ::tflite::ops::custom::Register_RANDOM_UNIFORM()); |
| #endif // TC3_WITH_ACTIONS_OPS |
| customize_fn(resolver.get()); |
| return std::unique_ptr<tflite::OpResolver>(std::move(resolver)); |
| } |
| |
| std::unique_ptr<const tflite::FlatBufferModel> TfLiteModelFromModelSpec( |
| const tflite::Model* model_spec) { |
| std::unique_ptr<const tflite::FlatBufferModel> model( |
| tflite::FlatBufferModel::BuildFromModel(model_spec)); |
| if (!model || !model->initialized()) { |
| TC3_LOG(ERROR) << "Could not build TFLite model from a model spec."; |
| return nullptr; |
| } |
| return model; |
| } |
| |
| std::unique_ptr<const tflite::FlatBufferModel> TfLiteModelFromBuffer( |
| const flatbuffers::Vector<uint8_t>* model_spec_buffer) { |
| const tflite::Model* model = |
| flatbuffers::GetRoot<tflite::Model>(model_spec_buffer->data()); |
| flatbuffers::Verifier verifier(model_spec_buffer->data(), |
| model_spec_buffer->size()); |
| if (!model->Verify(verifier)) { |
| return nullptr; |
| } |
| return TfLiteModelFromModelSpec(model); |
| } |
| |
| TfLiteModelExecutor::TfLiteModelExecutor( |
| std::unique_ptr<const tflite::FlatBufferModel> model) |
| : model_(std::move(model)), resolver_(BuildOpResolver()) {} |
| TfLiteModelExecutor::TfLiteModelExecutor( |
| std::unique_ptr<const tflite::FlatBufferModel> model, |
| std::unique_ptr<tflite::OpResolver> resolver) |
| : model_(std::move(model)), resolver_(std::move(resolver)) {} |
| |
| std::unique_ptr<tflite::Interpreter> TfLiteModelExecutor::CreateInterpreter() |
| const { |
| std::unique_ptr<tflite::Interpreter> interpreter; |
| tflite::InterpreterBuilder(*model_, *resolver_)(&interpreter); |
| return interpreter; |
| } |
| |
| template <> |
| void TfLiteModelExecutor::SetInput(const int input_index, |
| const std::vector<std::string>& input_data, |
| tflite::Interpreter* interpreter) const { |
| tflite::DynamicBuffer buf; |
| for (const std::string& s : input_data) { |
| buf.AddString(s.data(), s.length()); |
| } |
| buf.WriteToTensorAsVector( |
| interpreter->tensor(interpreter->inputs()[input_index])); |
| } |
| |
| template <> |
| std::vector<tflite::StringRef> TfLiteModelExecutor::Output( |
| const int output_index, const tflite::Interpreter* interpreter) const { |
| const TfLiteTensor* output_tensor = |
| interpreter->tensor(interpreter->outputs()[output_index]); |
| const int num_strings = tflite::GetStringCount(output_tensor); |
| std::vector<tflite::StringRef> output(num_strings); |
| for (int i = 0; i < num_strings; i++) { |
| output[i] = tflite::GetString(output_tensor, i); |
| } |
| return output; |
| } |
| |
| template <> |
| std::vector<std::string> TfLiteModelExecutor::Output( |
| const int output_index, const tflite::Interpreter* interpreter) const { |
| std::vector<std::string> output; |
| for (const tflite::StringRef& s : |
| Output<tflite::StringRef>(output_index, interpreter)) { |
| output.push_back(std::string(s.str, s.len)); |
| } |
| return output; |
| } |
| |
| } // namespace libtextclassifier3 |