litert_models

LiteRT (Formely TFLite) Model Tests

This test suite exercises LiteRT, formely known as TensorFlow Lite models. Most pretrained models are sourced from https://www.kaggle.com/models.

Testing currently follows several stages:

graph LR
  Model["Download model"]
  Model --> ImportMLIR["Import into MLIR"]
  ImportMLIR --> CompileIREE["Compile with IREE"]

Loading

Testing could also test inference and compare with LiteRT:

graph LR
  Model --> ImportMLIR["Import into MLIR"]
  ImportMLIR --> CompileIREE["Compile with IREE"]
  CompileIREE --> RunIREE["Run with IREE"]
  RunIREE --> Check

  Model --> LoadLiteRT["Load into LiteRT"]
  LoadLiteRT --> RunLiteRT["Run with LiteRT"]
  RunLiteRT --> Check

  Check["Compare results"]

Loading

Quickstart

1. Set up your virtual environment and install requirements:

   python -m venv .venv
   source .venv/bin/activate
   python -m pip install -r requirements.txt

   • To use iree-compile and iree-run-module from Python packages:

     python -m pip install -r requirements-iree.txt

   • To use a custom version of IREE follow the instructions for building the IREE Python packages from source, including the extra steps for the TFLite importer.

2. Run pytest using typical flags:

   pytest \
     -rA \
     --log-cli-level=info \
     --durations=0

   See https://docs.pytest.org/en/stable/how-to/usage.html for other options.

Advanced pytest usage

• The log-cli-level level can also be set to debug, warning, or error. See https://docs.pytest.org/en/stable/how-to/logging.html.

• Run only tests matching a name pattern:

  pytest -k resnet

• Ignore xfail marks (https://docs.pytest.org/en/stable/how-to/skipping.html#ignoring-xfail):

  pytest --runxfail

• Run tests in parallel using https://pytest-xdist.readthedocs.io/ (note that this swallows some logging):

  # Run with an automatic number of threads (usually one per CPU core).
  pytest -n auto
  
  # Run on an explicit number of threads.
  pytest -n 4

• Create an HTMl report using https://pytest-html.readthedocs.io/en/latest/index.html

  pytest --html=report.html --self-contained-html --log-cli-level=info

  See also https://docs.pytest.org/en/latest/how-to/output.html#creating-junitxml-format-files

Test suite implementation details

Kaggle

Models are downloaded using https://github.com/Kaggle/kagglehub.

By default, kagglehub caches downloads at ~/.cache/kagglehub/models/. This can be overriden by setting the KAGGLEHUB_CACHE environment variable. See the kagglehub/config.py source for other configuration options.

Working with .mlirbc files

The iree-import-tflite tool outputs MLIR bytecode (.mlirbc) by default. To convert to MLIR text (.mlir):

iree-ir-tool cp input.mlirbc -o output.mlir