From 55d84e75c29f45f9a277c01be46f19edef4e76c9 Mon Sep 17 00:00:00 2001 From: fban Date: Thu, 17 Nov 2022 22:59:58 +0000 Subject: 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 --- Android.bp | 2 - .../text/nlclassifier/BertNLClassifierTest.java | 1 - .../src/native/task/core/minimal_op_resolver.cc | 46 +++++++++++++++++----- 3 files changed, 37 insertions(+), 12 deletions(-) diff --git a/Android.bp b/Android.bp index 34cc018b..f3d5bb4e 100644 --- a/Android.bp +++ b/Android.bp @@ -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 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 -#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 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 CreateOpResolver() { ::tflite::ops::builtin::Register_PACK()); resolver.AddBuiltin(::tflite::BuiltinOperator_SOFTMAX, ::tflite::ops::builtin::Register_SOFTMAX()); - return absl::make_unique(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(resolver); } } // namespace task -- cgit v1.2.3