diff options
author | fban <fban@google.com> | 2022-11-17 22:59:58 +0000 |
---|---|---|
committer | Frank Ban <fban@google.com> | 2023-02-17 23:51:24 +0000 |
commit | 55d84e75c29f45f9a277c01be46f19edef4e76c9 (patch) | |
tree | 97cf973d8b5114e65bb9444dee1207cce6b36cd6 | |
parent | 4dd11643c33cecb69bfb1d3f06b6986ecf701023 (diff) | |
download | tflite-support-55d84e75c29f45f9a277c01be46f19edef4e76c9.tar.gz |
Adds a minimal op-resolver for Rubidium's usage.
This change adds a custom op-resolver that only includes the necessary ops for the
bert_nl_classifier and rb_model test models. This gives the Rb team the
opportunity for ~1M savings in their binary size.
For now, the actual classifier build target still uses the
builtin_op_resolver. Later, we can investigate how to set up native build targets
that can accept op-resolver dependencies at link-time for greater flexibility (like we use in Bazel).
Before and after (gpaste): 4845371938308096
Test: atest OnDeviceClassifierTest
Bug: 238435760
Change-Id: I6ca794fed5cf1cd3ffc0612a60754e4bbea9fc8f
3 files changed, 37 insertions, 12 deletions
@@ -225,7 +225,6 @@ cc_library_shared { "tensorflow_lite_support/java/src/native/task/core/builtin_op_resolver.cc", "tensorflow_lite_support/cc/utils/jni_utils.cc", ], - // TODO(b/247088924): Use linker_scripts here. version_script: "tensorflow_lite_support/java/tflite_version_script.lds", shared_libs: ["liblog"], static_libs: [ @@ -246,7 +245,6 @@ cc_library_shared { "tensorflow_headers", "flatbuffer_headers", "jni_headers", - "liblog_headers", "libtextclassifier_flatbuffer_headers", ], generated_headers: [ diff --git a/tensorflow_lite_support/java/src/javatests/org/tensorflow/lite/task/text/nlclassifier/BertNLClassifierTest.java b/tensorflow_lite_support/java/src/javatests/org/tensorflow/lite/task/text/nlclassifier/BertNLClassifierTest.java index 8c71f705..efaa9d99 100644 --- a/tensorflow_lite_support/java/src/javatests/org/tensorflow/lite/task/text/nlclassifier/BertNLClassifierTest.java +++ b/tensorflow_lite_support/java/src/javatests/org/tensorflow/lite/task/text/nlclassifier/BertNLClassifierTest.java @@ -124,6 +124,5 @@ public class BertNLClassifierTest { private void verifyDynamicInputResults(BertNLClassifier classifier) { List<Category> topics = classifier.classify("FooBarBaz"); assertThat(topics.size()).isEqualTo(446); - // TODO(ag/19888344): Add a test for a long text input. } } diff --git a/tensorflow_lite_support/java/src/native/task/core/minimal_op_resolver.cc b/tensorflow_lite_support/java/src/native/task/core/minimal_op_resolver.cc index 31d693a8..32d1054d 100644 --- a/tensorflow_lite_support/java/src/native/task/core/minimal_op_resolver.cc +++ b/tensorflow_lite_support/java/src/native/task/core/minimal_op_resolver.cc @@ -15,7 +15,6 @@ limitations under the License. #include <memory> -#include "absl/memory/memory.h" #include "tensorflow/lite/kernels/builtin_op_kernels.h" #include "tensorflow/lite/op_resolver.h" @@ -23,21 +22,15 @@ namespace tflite { namespace task { // Create a minimal MutableOpResolver to provide only -// the ops required by NLClassifier/BertNLClassifier. +// the ops required by the bert_nl_classifier and rb_model for BertNLClassifier. std::unique_ptr<MutableOpResolver> CreateOpResolver() { MutableOpResolver resolver; - resolver.AddBuiltin(::tflite::BuiltinOperator_DEQUANTIZE, - ::tflite::ops::builtin::Register_DEQUANTIZE()); resolver.AddBuiltin(::tflite::BuiltinOperator_RESHAPE, ::tflite::ops::builtin::Register_RESHAPE()); resolver.AddBuiltin(::tflite::BuiltinOperator_GATHER, ::tflite::ops::builtin::Register_GATHER()); resolver.AddBuiltin(::tflite::BuiltinOperator_STRIDED_SLICE, ::tflite::ops::builtin::Register_STRIDED_SLICE()); - resolver.AddBuiltin(::tflite::BuiltinOperator_PAD, - ::tflite::ops::builtin::Register_PAD()); - resolver.AddBuiltin(::tflite::BuiltinOperator_CONCATENATION, - ::tflite::ops::builtin::Register_CONCATENATION()); resolver.AddBuiltin(::tflite::BuiltinOperator_FULLY_CONNECTED, ::tflite::ops::builtin::Register_FULLY_CONNECTED()); resolver.AddBuiltin(::tflite::BuiltinOperator_CAST, @@ -54,7 +47,42 @@ std::unique_ptr<MutableOpResolver> CreateOpResolver() { ::tflite::ops::builtin::Register_PACK()); resolver.AddBuiltin(::tflite::BuiltinOperator_SOFTMAX, ::tflite::ops::builtin::Register_SOFTMAX()); - return absl::make_unique<MutableOpResolver>(resolver); + resolver.AddBuiltin(::tflite::BuiltinOperator_EXPAND_DIMS, + ::tflite::ops::builtin::Register_EXPAND_DIMS()); + resolver.AddBuiltin(::tflite::BuiltinOperator_SHAPE, + ::tflite::ops::builtin::Register_SHAPE()); + resolver.AddBuiltin(::tflite::BuiltinOperator_FILL, + ::tflite::ops::builtin::Register_FILL()); + resolver.AddBuiltin(::tflite::BuiltinOperator_SUB, + ::tflite::ops::builtin::Register_SUB()); + resolver.AddBuiltin(::tflite::BuiltinOperator_MEAN, + ::tflite::ops::builtin::Register_MEAN()); + 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_BATCH_MATMUL, + ::tflite::ops::builtin::Register_BATCH_MATMUL()); + resolver.AddBuiltin(::tflite::BuiltinOperator_GELU, + ::tflite::ops::builtin::Register_GELU()); + resolver.AddBuiltin(::tflite::BuiltinOperator_TANH, + ::tflite::ops::builtin::Register_TANH()); + resolver.AddBuiltin(::tflite::BuiltinOperator_LOGISTIC, + ::tflite::ops::builtin::Register_LOGISTIC()); + resolver.AddBuiltin(::tflite::BuiltinOperator_SLICE, + ::tflite::ops::builtin::Register_SLICE()); + // Needed for the test bert_nl_classifier model. + resolver.AddBuiltin(::tflite::BuiltinOperator_PAD, + ::tflite::ops::builtin::Register_PAD()); + resolver.AddBuiltin(::tflite::BuiltinOperator_CONCATENATION, + ::tflite::ops::builtin::Register_CONCATENATION()); + resolver.AddBuiltin(::tflite::BuiltinOperator_FULLY_CONNECTED, + ::tflite::ops::builtin::Register_FULLY_CONNECTED(), + /*version=*/9); + resolver.AddBuiltin(::tflite::BuiltinOperator_DEQUANTIZE, + ::tflite::ops::builtin::Register_DEQUANTIZE(), + /*version=*/2); + return std::make_unique<MutableOpResolver>(resolver); } } // namespace task |