Native Training Example

This example shows how to

1. Build a PyTorch functional model for training
2. Import that model into IREE's compiler
3. Compile that model to an IREE VM bytecode module
4. Load the compiled module using IREE's high level runtime C API into a lightweight program
5. Train the loaded model

This example was built with the goal of allowing you to be able to build it outside this repo in your own project with minimal changes.

The weights for the model are stored in the program itself and updated in memory. This can be modified to be stored however you see fit.

Running the Example

Install iree-torch and other dependencies necessary for this example. iree-torch provides a number of convenient wrappers around torch-mlir and iree compilation:

Note We recommend installing Python packages inside a virtual environment.

pip install -f https://iree-org.github.io/iree/pip-release-links.html iree-compiler
pip install -f https://llvm.github.io/torch-mlir/package-index/ torch-mlir
pip install git+https://github.com/iree-org/iree-torch.git

Update submodules in this repo:

(cd $(git rev-parse --show-toplevel) && git submodule update --init)

Build the IREE runtime:

(cd $(git rev-parse --show-toplevel) && cmake -GNinja -B /tmp/iree-build-runtime/ .)
cmake --build /tmp/iree-build-runtime/ --target iree_runtime_unified

Make sure you‘re in this example’s directory:

cd $(git rev-parse --show-toplevel)/samples/native_training

Build the native training example:

make

Generate the IREE VM bytecode for the model:

python native_training.py /tmp/native_training.vmfb

Run the native training model:

./native-training /tmp/native_training.vmfb