Skip to content

TensorFlow integrationlink

Overviewlink

IREE supports compiling and running TensorFlow programs represented as tf.Module classes or stored in the SavedModel format.

graph LR
  accTitle: TensorFlow to runtime deployment workflow overview
  accDescr {
    Programs start as either TensorFlow SavedModel or tf.Module programs.
    Programs are imported into MLIR as StableHLO.
    The IREE compiler uses the imported MLIR.
    Compiled programs are used by the runtime.
  }

  subgraph A[TensorFlow]
    direction TB
    A1[SavedModel]
    A2[tf.Module]

    A1 --- A2
  end

  subgraph B[MLIR]
    B1[StableHLO]
  end

  C[IREE compiler]
  D[Runtime deployment]

  A -- iree-import-tf --> B
  B --> C
  C --> D

Prerequisiteslink

  1. Install TensorFlow by following the official documentation:

    python -m pip install tf-nightly

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

    Stable release packages are published to PyPI.

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

    Nightly releases are published on GitHub releases.

    python -m pip install \
  --find-links https://iree.dev/pip-release-links.html \
  --upgrade \
  iree-compiler \
  iree-runtime \
  iree-tools-tf

Importing modelslink

IREE compilers transform a model into its final deployable format in several sequential steps. The first step for a TensorFlow model is to use either the iree-import-tf command-line tool or IREE's Python APIs to import the model into a format (i.e., MLIR) compatible with the generic IREE compilers.

From SavedModel on TensorFlow Hublink

IREE supports importing and using SavedModels from TensorFlow Hub.

Using the command-line toollink

First download the SavedModel and load it to get the serving signature, which is used as the entry point for IREE compilation flow:

import tensorflow.compat.v2 as tf
loaded_model = tf.saved_model.load('/path/to/downloaded/model/')
print(list(loaded_model.signatures.keys()))

Note

If there are no serving signatures in the original SavedModel, you may add them by yourself by following "Missing serving signature in SavedModel".

Then you can import the model with iree-import-tf. You can read the options supported via iree-import-tf -help. Using MobileNet v2 as an example and assuming the serving signature is predict:

iree-import-tf
  --tf-import-type=savedmodel_v1 \
  --tf-savedmodel-exported-names=predict \
  /path/to/savedmodel -o iree_input.mlir

Tip

iree-import-tf is installed as /path/to/python/site-packages/iree/tools/tf/iree-import-tf. You can find out the full path to the site-packages directory via the python -m site command.

Tip

-tf-import-type needs to match the SavedModel version. You can try both v1 and v2 if you see one of them gives an empty dump.

Next, you can compile the model in iree_input.mlir for one of IREE's supported targets by following one of the deployment configuration guides.

Sampleslink

Colab notebooks
Training an MNIST digits classifier Open in Colab
Edge detection Open In Colab
Pretrained ResNet50 inference Open In Colab
TensorFlow Hub import Open In Colab

End-to-end execution tests can be found in IREE's integrations/tensorflow/e2e/ directory.

Troubleshootinglink

Missing serving signature in SavedModellink

Sometimes SavedModels are exported without explicit serving signatures. This happens by default for TensorFlow Hub SavedModels. However, serving signatures are required as entry points for IREE compilation flow. You can use Python to load and re-export the SavedModel to give it serving signatures. For example, for MobileNet v2, assuming we want the serving signature to be predict and operating on a 224x224 RGB image:

import tensorflow.compat.v2 as tf
loaded_model = tf.saved_model.load('/path/to/downloaded/model/')
call = loaded_model.__call__.get_concrete_function(
         tf.TensorSpec([1, 224, 224, 3], tf.float32))
signatures = {'predict': call}
tf.saved_model.save(loaded_model,
  '/path/to/resaved/model/', signatures=signatures)

The above will create a new SavedModel with a serving signature, predict, and save it to /path/to/resaved/model/.