/*
* 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"
// Forward declaration of custom TensorFlow Lite ops for registration.
namespace tflite {
namespace ops {
namespace builtin {
TfLiteRegistration* Register_ADD();
TfLiteRegistration* Register_CONCATENATION();
TfLiteRegistration* Register_CONV_2D();
TfLiteRegistration* Register_FULLY_CONNECTED();
TfLiteRegistration* Register_L2_NORMALIZATION();
TfLiteRegistration* Register_MUL();
TfLiteRegistration* Register_RESHAPE();
TfLiteRegistration* Register_SOFTMAX();
TfLiteRegistration* Register_GATHER();
TfLiteRegistration* Register_TRANSPOSE();
TfLiteRegistration* Register_SUB();
TfLiteRegistration* Register_DIV();
TfLiteRegistration* Register_STRIDED_SLICE();
TfLiteRegistration* Register_EXP();
TfLiteRegistration* Register_TOPK_V2();
TfLiteRegistration* Register_SPLIT();
TfLiteRegistration* Register_CAST();
TfLiteRegistration* Register_MAXIMUM();
TfLiteRegistration* Register_MINIMUM();
TfLiteRegistration* Register_NEG();
TfLiteRegistration* Register_SLICE();
TfLiteRegistration* Register_LOG();
TfLiteRegistration* Register_SUM();
TfLiteRegistration* Register_PACK();
TfLiteRegistration* Register_DEQUANTIZE();
TfLiteRegistration* Register_MEAN();
} // namespace builtin
} // namespace ops
} // namespace tflite
#ifdef TC3_WITH_ACTIONS_OPS
#include "utils/tflite/dist_diversification.h"
#include "utils/tflite/text_encoder.h"
#include "utils/tflite/token_encoder.h"
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=*/3);
resolver->AddBuiltin(tflite::BuiltinOperator_FULLY_CONNECTED,
tflite::ops::builtin::Register_FULLY_CONNECTED(),
/*min_version=*/1,
/*max_version=*/4);
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_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_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_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_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());
}
#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 {
inline std::unique_ptr<tflite::OpResolver> BuildOpResolver() {
#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("TokenEncoder",
tflite::ops::custom::Register_TOKEN_ENCODER());
#endif // TC3_WITH_ACTIONS_OPS
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->Length());
if (!model->Verify(verifier)) {
return nullptr;
}
return TfLiteModelFromModelSpec(model);
}
TfLiteModelExecutor::TfLiteModelExecutor(
std::unique_ptr<const tflite::FlatBufferModel> model)
: model_(std::move(model)), resolver_(BuildOpResolver()) {}
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