Quickstart-android: Running TFLite model with custom operations

[the-dagger]

Hi!
Can Firebase MLKit run TfLite models with custom operations?

The provided examples use models with built in operations but not models with custom ops.
For example, the smartreply model has 3 custom operations which gives an error while loading the Model using MLKit.

Stacktrace :

 W/System.err: Caused by: com.google.firebase.ml.common.FirebaseMLException: Internal error has occurred when executing Firebase ML tasks
        at com.google.android.gms.internal.firebase_ml.zzgp.zza(Unknown Source:15)
        ... 5 more
    Caused by: java.lang.IllegalArgumentException: Cannot create interpreter: Didn't find custom op for name 'Normalize'
07-23 18:01:16.018 21726-21726/io.github.the_dagger.mlkit W/System.err: Didn't find custom op for name 'ExtractFeatures'
    Didn't find custom op for name 'Predict'
    Registration failed.
        at org.tensorflow.lite.NativeInterpreterWrapper.createInterpreter(Native Method)
        at org.tensorflow.lite.NativeInterpreterWrapper.<init>(NativeInterpreterWrapper.java:63)
        at org.tensorflow.lite.NativeInterpreterWrapper.<init>(NativeInterpreterWrapper.java:51)
        at org.tensorflow.lite.Interpreter.<init>(Interpreter.java:90)
07-23 18:01:16.019 21726-21726/io.github.the_dagger.mlkit W/System.err:     at com.google.android.gms.internal.firebase_ml.zzhn.zzey(Unknown Source:102)
        at com.google.android.gms.internal.firebase_ml.zzha.zzfb(Unknown Source:75)
        at com.google.android.gms.internal.firebase_ml.zzha.call(Unknown Source:0)
        at com.google.android.gms.internal.firebase_ml.zzgp.zza(Unknown Source:0)
07-23 18:01:16.020 21726-21726/io.github.the_dagger.mlkit W/System.err:     ... 5 more
07-23 18:01:21.529 21726-21733/io.github.the_dagger.mlkit I/zygote64: Do partial code cache collection, code=123KB, data=88KB
07-23 18:01:21.530 21726-21733/io.github.the_dagger.mlkit I/zygote64: After code cache collection, code=123KB, data=88KB
    Increasing code cache capacity to 512KB

The relevant code can be found here :
https://github.com/the-dagger/MLKitAndroid/blob/master/app/src/main/java/io/github/the_dagger/mlkit/activity/SmartReplyActivity.kt

P.S. Docs here outline that I should write native code for custom operations to be used with TFLite but I couldn't find a way on doing it with ML Kit.

https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/lite/g3doc/custom_operators.md

Thanks!

[samtstern]

cc @isaisachen @gkaldev does either of you know about custom operations with ML Kit?

[the-dagger]

@samtstern Managed to resolve this issue by building a custom tflite build with source code for the missing custom operations (Yes, I had a fun filled weekend).

Using a custom build resolved this issue, but presented me with another issue which I am clueless about.

07-29 00:14:07.340 6489-6489/? W/System.err:     at com.google.android.gms.internal.firebase_ml.zzgq.run(Unknown Source:4)
        at android.os.Handler.handleCallback(Handler.java:873)
        at android.os.Handler.dispatchMessage(Handler.java:99)
        at android.os.Looper.loop(Looper.java:193)
        at android.os.HandlerThread.run(HandlerThread.java:65)
    Caused by: java.lang.IllegalArgumentException: DataType error: DataType 5 is not recognized in Java (version 3)
        at org.tensorflow.lite.DataType.fromNumber(DataType.java:54)
        at org.tensorflow.lite.Tensor.<init>(Tensor.java:211)
        at org.tensorflow.lite.Tensor.fromHandle(Tensor.java:32)
        at org.tensorflow.lite.NativeInterpreterWrapper.getInputTensor(NativeInterpreterWrapper.java:245)
        at org.tensorflow.lite.NativeInterpreterWrapper.run(NativeInterpreterWrapper.java:110)
        at org.tensorflow.lite.Interpreter.runForMultipleInputsOutputs(Interpreter.java:168)
        at com.google.android.gms.internal.firebase_ml.zzho.runForMultipleInputsOutputs(Unknown Source:2)
        at com.google.android.gms.internal.firebase_ml.zzhn.zza(Unknown Source:459)
        at com.google.android.gms.internal.firebase_ml.zzhn.zza(Unknown Source:2)
        at com.google.android.gms.internal.firebase_ml.zzgu.call(Unknown Source:15)
        at com.google.android.gms.internal.firebase_ml.zzgp.zza(Unknown Source:0)
        ... 5 more

I'm using the smartreply tfLite model which expects input as a String and returns the output as String, but the default datatypes in TFLite only support FLOAT,INT,BYTE and LONG.

The default example over in the docs uses image detection which works well with INPUT_TYPE set to BYTE and OUTPUT_TYPE set to Int, but there's no explanation whatsoever for String inputs/outputs.

@isaisachen @gkaldev any leads would be appreciated!

[samtstern]

@the-dagger thank you for the update! That's super cool that you managed to get a custom build going, let's see if I can find someone about that DataType 5 issue.

[psanketi]

@the-dagger Thanks. Awesome that you got to build your own custom TFLite, which is needed if you want to have custom TFLite Ops.
You are right that String is not natively supported in MLKit at the moment. We should be able to add that. Let me try to get back to if we can find a workaround in the mean time using existing types.

[the-dagger]

@psanketi @samtstern Thanks for the update!
Dince on exploring further it turns out that even the TfLite data type only supports 4 data types, I opened a similar issue on TFLite github repo as well.
Turns out that it currently doesn't support dynamic/static String input and output and this is something that's on the team's radar.

@psanketi Would it be possible to hack around and modify the current custom op that accepts a String input to accept a Byte array instead?
I tried this but that ended up with the same DataType 5 error.

[isaisachen]

@the-dagger Thanks for reporting.

As @psanketi and @jdduke responded earlier, String is not well supported in Java APIs. It's only supported with C++ API.

Could you describe your planned use case? As of now, if you want to keep using Java APIs (ML Kit or TF Lite), you will have to modify ops to do data parsing and normalizing yourself:
1) The open-sourced smart reply implementation uses string_utils here which populates a String Tensor w/ pre-defined protocol.
2) ML Kit is expecting fixed output dimension which is not the case for smart reply.

All above doesn't sound like worth the effort.

If you just want to model against smart reply and implement some new custom model yourself, I'd suggest using TFLite C++ API directly. Smart Reply TFLite integration code is here

[the-dagger]

Hi @isaisachen thanks for the detailed explanation, it makes sense now.
As for the user case, I was just playing around with the given models in an attempt to implement the upcoming MLKit API for smart reply myself using the smart reply model.

I can do it using the C++ APIs, but I wanted to use it with Firebase MLKit for a blog I was planning to write on the same.

But anyways, I have decided to use a different model for the blog for now, will wait for the Java APIs to support String input/output.

@psanketi Do let me know in case you stumble upon any workaround for this using the Java APIs and built in APIs.

Thanks!

[isaisachen]

@the-dagger I see. Cool. Could you shoot a quick email to @samtstern? It's listed on his profile. Thanks.

[samtstern]

@the-dagger just so you know, your comments have spawned some really productive internal conversations about how we can support use cases like yours in the future. No concrete plans, just want to say thanks for pushing us in the right direction.

[the-dagger]

@samtstern @isaisachen you're welcome, glad to be useful.

[samtstern]

We've got some feature requests and bugs filed internally for this discussion. Going to close it here as there's nothing we can do in this repo in the short/medium term. Will update if I hear anything!

[the-dagger]

Since I couldn't find a good guide that outlined building a TensorFlow Lite build with custom Operations, I went ahead wrote one.

https://heartbeat.fritz.ai/compiling-a-tensorflow-lite-build-with-custom-operations-cf6330ee30e2

Will probably help others who stumble upon this issue :slightly_smiling_face: