Skip to content

TensorFlow Lite Integration

IREE supports compiling and running pre-trained TensorFlow Lite (TFLite) models. It converts a model to TOSA MLIR, then compiles it into a VM module.

Prerequisites

Download a pre-trained TFLite model from the list of hosted models, or use the TensorFlow Lite converter to convert a TensorFlow model to a .tflite flatbuffer.

Install IREE pip packages, either from pip or by building from source:

python -m pip install \
  iree-compiler \
  iree-runtime \
  iree-tools-tflite

Warning

The TensorFlow Lite package is currently only available on Linux and macOS. It is not available on Windows yet (see this issue).

Importing models

Fist, import the TFLite model to TOSA MLIR:

iree-import-tflite \
  sample.tflite \
  -o sample.mlir

Next, compile the TOSA MLIR to a VM flatbuffer, using either the command line tools or the Python API:

Using the command-line tool

iree-translate \
  --iree-mlir-to-vm-bytecode-module \
  --iree-input-type=tosa \
  --iree-hal-target-backends=vmvx \
  sample.mlir \
  -o sample.vmfb

Using the python API

from iree.compiler import compile_str
with open('sample.mlir') as sample_tosa_mlir:
  compiled_flatbuffer = compile_str(sample_tosa_mlir.read(),
    input_type="tosa",
    target_backends=["vmvx"],
    extra_args=["--iree-native-bindings-support=false",
      "--iree-tflite-bindings-support"])

Todo

Issue#5462: Link to TensorFlow Lite bindings documentation once it has been written.

The flatbuffer can then be loaded to a VM module and run through IREE's runtime.

Samples

Colab notebooks
Text classification with TFLite and IREE Open in Colab

Todo

Issue#3954: Add documentation for an Android demo using the Java TFLite bindings, once it is complete at not-jenni/iree-android-tflite-demo.

Back to top