LiteRT Next crashes during inference

[aicelen]

Hey!
I want to add LiteRT support to the Python Kivy framework. Therefore I wrote this .
A small image classifier is working perfectly fine. A bigger image segmentation network is causing the app to crash. This only occurs on android when using LiteRT Next (not on wsl or android with LiteRT). The model also works when using the tflite benchmark app.
All the code used can be found in the linked repo.

Logs:

07-08 21:29:32.274 11747 11789 D nativeloader: Load /data/app/~~O19ptvLYjv0kWcpdRMlsgA==/org.test.rt_next-f0P71dabFO12TjeaUA6BEA==/lib/arm64/liblitert_jni.so using class loader ns clns-7 (caller=/data/app/~~O19ptvLYjv0kWcpdRMlsgA==/org.test.rt_next-f0P71dabFO12TjeaUA6BEA==/base.apk): ok
07-08 21:29:32.303 11747 11789 I tflite  : [third_party/odml/litert/litert/core/environment.cc:26] Creating LiteRT environment with options
07-08 21:29:32.303 11747 11789 E tflite  : [third_party/odml/litert/litert/runtime/accelerators/dispatch/dispatch_accelerator.cc:113] Dispatch library directory is not set.
07-08 21:29:32.303 11747 11789 W tflite  : [third_party/odml/litert/litert/runtime/accelerators/auto_registration.cc:46] NPU accelerator could not be loaded and registered: kLiteRtStatusErrorInvalidArgument.
07-08 21:29:32.304 11747 11789 I tflite  : [third_party/odml/litert/litert/runtime/accelerators/auto_registration.cc:54] Dynamically loaded GPU accelerator registered.
07-08 21:29:32.304 11747 11789 I tflite  : [third_party/odml/litert/litert/runtime/accelerators/auto_registration.cc:70] CPU accelerator registered.
07-08 21:29:32.305 11747 11789 I tflite  : [third_party/odml/litert/litert/runtime/compiled_model.cc:121] Applying compiler plugins...
07-08 21:29:32.305 11747 11789 W tflite  : [third_party/odml/litert/litert/runtime/compiled_model.cc:129] Failed to apply compiler plugins: Compiler plugin is not configured
07-08 21:29:32.305 11747 11789 I tflite  : [third_party/odml/litert/litert/runtime/compiled_model.cc:148] Flatbuffer model initialized directly from incoming litert model.
07-08 21:29:32.322 11747 11789 I tflite  : [third_party/odml/litert/litert/runtime/compiled_model.cc:391] Tensor input_tensor is shared with CPU.
07-08 21:29:34.108 11798 11798 F DEBUG   :       #12 pc 000000000005a11c  /data/app/~~O19ptvLYjv0kWcpdRMlsgA==/org.test.rt_next-f0P71dabFO12TjeaUA6BEA==/lib/arm64/liblitert_jni.so (BuildId: f8c03dcfd5da8677b7fcbf9cb1bb1cbe)
07-08 21:29:34.108 11798 11798 F DEBUG   :       #13 pc 000000000005a22c  /data/app/~~O19ptvLYjv0kWcpdRMlsgA==/org.test.rt_next-f0P71dabFO12TjeaUA6BEA==/lib/arm64/liblitert_jni.so (BuildId: f8c03dcfd5da8677b7fcbf9cb1bb1cbe)
07-08 21:29:34.108 11798 11798 F DEBUG   :       #14 pc 0000000000049a30  /data/app/~~O19ptvLYjv0kWcpdRMlsgA==/org.test.rt_next-f0P71dabFO12TjeaUA6BEA==/lib/arm64/liblitert_jni.so (BuildId: f8c03dcfd5da8677b7fcbf9cb1bb1cbe)
07-08 21:29:34.108 11798 11798 F DEBUG   :       #15 pc 0000000000049888  /data/app/~~O19ptvLYjv0kWcpdRMlsgA==/org.test.rt_next-f0P71dabFO12TjeaUA6BEA==/lib/arm64/liblitert_jni.so (Java_com_google_ai_edge_litert_CompiledModel_nativeRun+324) (BuildId: f8c03dcfd5da8677b7fcbf9cb1bb1cbe)
07-08 21:29:34.108 11798 11798 F DEBUG   :       #22 pc 00000000001f4ec0  [anon:dalvik-classes.dex extracted in memory from /data/app/~~O19ptvLYjv0kWcpdRMlsgA==/org.test.rt_next-f0P71dabFO12TjeaUA6BEA==/base.apk] (com.google.ai.edge.litert.CompiledModel.access$nativeRun+0)
07-08 21:29:34.108 11798 11798 F DEBUG   :       #27 pc 00000000001f36c4  [anon:dalvik-classes.dex extracted in memory from /data/app/~~O19ptvLYjv0kWcpdRMlsgA==/org.test.rt_next-f0P71dabFO12TjeaUA6BEA==/base.apk] (com.google.ai.edge.litert.CompiledModel$Companion.nativeRun+0)
07-08 21:29:34.108 11798 11798 F DEBUG   :       #32 pc 00000000001f3664  [anon:dalvik-classes.dex extracted in memory from /data/app/~~O19ptvLYjv0kWcpdRMlsgA==/org.test.rt_next-f0P71dabFO12TjeaUA6BEA==/base.apk] (com.google.ai.edge.litert.CompiledModel$Companion.access$nativeRun+0)
07-08 21:29:34.108 11798 11798 F DEBUG   :       #37 pc 00000000001f4fd0  [anon:dalvik-classes.dex extracted in memory from /data/app/~~O19ptvLYjv0kWcpdRMlsgA==/org.test.rt_next-f0P71dabFO12TjeaUA6BEA==/base.apk] (com.google.ai.edge.litert.CompiledModel.run+0)
07-08 21:29:34.108 11798 11798 F DEBUG   :       #42 pc 00000000001f4f50  [anon:dalvik-classes.dex extracted in memory from /data/app/~~O19ptvLYjv0kWcpdRMlsgA==/org.test.rt_next-f0P71dabFO12TjeaUA6BEA==/base.apk] (com.google.ai.edge.litert.CompiledModel.run$default+0)
07-08 21:29:34.108 11798 11798 F DEBUG   :       #47 pc 00000000001f4f98  [anon:dalvik-classes.dex extracted in memory from /data/app/~~O19ptvLYjv0kWcpdRMlsgA==/org.test.rt_next-f0P71dabFO12TjeaUA6BEA==/base.apk] (com.google.ai.edge.litert.CompiledModel.run+0)

Note that the segmentation network is converted to .tflite using this code:

import tensorflow as tf
from keras.models import model_from_json

name = 'arch'

arch_path = f'/content/{name}.json'
model_path = f'/content/{name}.h5'

with open(arch_path, 'r') as f:
    model = model_from_json(f.read())

model.load_weights(model_path)

# Create a concrete function with fixed batch size
input_spec = tf.TensorSpec(shape=[1, 256, 256, 3], dtype=tf.float32)
@tf.function(input_signature=[input_spec])
def model_fixed(input_tensor):
    return model(input_tensor)

concrete_func = model_fixed.get_concrete_function()


# Convert with explicit shapes
converter = tf.lite.TFLiteConverter.from_concrete_functions([concrete_func])
converter.optimizations = [tf.lite.Optimize.DEFAULT]


tflite_model = converter.convert()

with open('unet_big_fixed.tflite', 'wb') as f:
    f.write(tflite_model)

Interestingly it takes around 1.5s until the app crashes. So the model is failing during inference.
Also I get logs until it reaches model.run(input, output).

I did not found a LiteRT Next benchmark app nor official documentation of model copatibility.

Thank you in advance,
aicelen

[aicelen]

The problem is the dynamic range quantization which is currently unsupported/not working. Without converter.optimizations = [tf.lite.Optimize.DEFAULT] the model is running fine.
I'll leave this open because it is a bug or requested_feature.

[gaikwadrahul8]

Hi, @aicelen
I apologize for the delay in my response, The LiteRT Next APIs are not compatible with the LiteRT APIs, so applications using LiteRT must completely migrate to LiteRT Next in order to make use of the features and capabilities provided by the new APIs. Applications cannot use the TF Lite Interpreter APIs and Compiled Model APIs interchangeably.

LiteRT Next provides APIs for Kotlin and C++. Applications using a LiteRT SDK in other languages should continue using LiteRT. I believe you're referring this official documentation and LiteRT Next is an alpha release and under active development. Note: LiteRT NPU acceleration is only available through an Early Access Program. If you are not already enrolled, Sign Up. Please refer this official documentation

I have no access to your model at the moment so there might be some unsupported ops causing this crash so you can give it try with post-training Full Integer Quantization (int8) or If full integer quantization proves too challenging (e.g. significant accuracy drop, difficulty with representative dataset), float16 is a good next step. It still offers benefits over FP32 and might be more compatible with the NPU than dynamic range quantization. Please refer this official documentation

If issue still persists if possible could you please help us with your Github repo along with model to reproduce the same behavior from our end to investigate this issue further from our end. Thank you for your cooperation and patience

[aicelen]

Thank you for your response. I already tried it with FP16 which is working. I don't think that this is a bug related to my code since I'm using Kotlin APIs in Python. I'll give it a try with a simple tflite quantized with dynamic int8. Shall I close this issue and open a new one if the problems persist (so it's easier to understand)?