Convert with a subset of TFLite built-in OPs

[hayhan]

Description of the bug:

The tensorflow lite builtin ops have grown to 209 (defined in tensorflow/lite/builtin_ops.h). While some devices may still run old tflite runtime such as tf 2.14.1 which cannot be upgraded for some reasons. The currently converted tflite model with ai-edge-torch may include newer OPs that the old tflite runtime doesn't support.

Is it possible to run the converter with a subset of OPs, e.g. with the ones of (0 ~ 161)?

Actual vs expected behavior:

No response

Any other information you'd like to share?

No response

[pkgoogle]

Hi @hayhan, can you try the below and let me know if it satisfies your issues?

import torch
import torchvision
import ai_edge_torch
import tensorflow as tf

# 1. Load the PyTorch model
resnet18 = torchvision.models.resnet18(
    torchvision.models.ResNet18_Weights.IMAGENET1K_V1
).eval()
sample_inputs = (torch.randn(1, 3, 224, 224),)

# 2. Define the exact set of TFLite operators to allow.
# For compatibility with older runtimes, the user should specify only the
# operator sets that their runtime supports. TFLITE_BUILTINS is the most basic set.
allowed_op_sets = [
    tf.lite.OpsSet.TFLITE_BUILTINS
]

# 3. Create the converter flags dictionary.
# This dictionary will be passed to the TFLite converter.
tfl_converter_flags = {
    "target_spec": {
        "supported_ops": allowed_op_sets
    }
}

# 4. Convert the model using the flags.
# The `_ai_edge_converter_flags` argument passes the configuration.
edge_model = ai_edge_torch.convert(
    resnet18,
    sample_inputs,
    _ai_edge_converter_flags=tfl_converter_flags,
)

# 5. Export the model
edge_model.export("resnet18_restricted_ops.tflite")

[hayhan]

Hi @pkgoogle, I am trying to convert the smollm2-135m-instruct model in the generative examples folder. I followed your instruction and gave "edge_model = converter.convert()" an extra arg you listed above. But the generated tflite still contains the op: StablehloComposite (#206 in tensorflow/lite/builtin_ops.h).

I suppose the tf.lite.OpsSet.TFLITE_BUILTINS means ALL OPs in tensorflow/lite/builtin_ops.h, right? How can I convert without the StablehloComposite (#206)? The TFLite runtime (2.14.1) on my device doesn't support this OP. Thanks.

[Hardik-Choraria]

Hi @pkgoogle,

I’m also trying to convert the qwen2-0.5B model in the generative_examples folder. I followed your instructions and passed the extra argument to converter.convert() but the generated TFLite still contain the op StablehloScatter.

[pkgoogle]

Hi @hayhan, @Hardik-Choraria, any way you can update the runtime on your device? generally speaking the later versions and especially LiteRT will be more stable. I cannot determine an easy (non-code change way) to make this work easily... perhaps changing the tensorflow submodule to be a version before those ops were added to BUILT-INs that might work but I think there will be a lot of version mismatch issues afterwards. Otherwise I'm thinking we would need to create an "Device -> supported Ops Map registry" to help with cases like this.

[github-actions]

Marking this issue as stale since it has been open for 7 days with no activity. This issue will be closed if no further activity occurs.

[github-actions]

This issue was closed because it has been inactive for 14 days. Please post a new issue if you need further assistance. Thanks!