Treelite : toolbox for decision tree deployment

Treelite is a flexible toolbox for efficient deployment of decision tree ensembles.

Star Watch

You are currently browsing the documentation of a development version of treelite: commit bd5529a.

Why treelite?

Use machine learning package of your choice

Treelite accommodates a wide range of decision tree ensemble models. In particular, it handles both random forests and gradient boosted trees.

Treelite can read models produced by XGBoost, LightGBM, and scikit-learn. In cases where you are using another package to train your model, you may use the flexible builder class.

Deploy with minimal dependencies

It is a great hassle to install machine learning packages (e.g. XGBoost, LightGBM, scikit-learn, etc.) on every machine your tree model will run. This is the case no longer: treelite will export your model as a stand-alone prediction subroutine so that predictions will be made without any machine learning package installed.

Compile and optimize your model for fast prediction

Treelite optimizes the prediction subroutine for faster prediction.

Depending on your use cases, simply compiling the prediction subroutine into machine code may boost the performance noticeably. (See the benchmark section below.) In addition, treelite supports additional optimizations that improve performance while preserving the ensemble model.

How treelite works


(Click to enlarge)

The workflow involves two distinct machines: the host machine that generates prediction subroutine from a given tree model, and the target machine that runs the subroutine. The two machines exchange a single C file that contains all relevant information about the tree model. Only the host machine needs to have treelite installed; the target machine requires only a working C compiler.

Quick start

Install treelite from PyPI:

pip3 install --user treelite

Import your tree ensemble model into treelite:

import treelite
model = treelite.Model.load('my_model.model', format='xgboost')

Deploy a source archive:

# Produce a zipped source directory, containing all model information
# Run `make` on the target machine
model.export_srcpkg(platform='unix', toolchain='gcc',
                    pkgpath='./mymodel.zip', libname='mymodel.so',
                    verbose=True)

Deploy a shared library:

# Like export_srcpkg, but generates a shared library immediately
# Use this only when the host and target machines are compatible
model.export_lib(toolchain='gcc', libpath='./mymodel.so', verbose=True)

Make predictions on the target machine:

import treelite.runtime
predictor = treelite.runtime.Predictor('./mymodel.so', verbose=True)
batch = treelite.runtime.Batch.from_npy2d(X)
out_pred = predictor.predict(batch, verbose=True)

Read First tutorial for a more detailed example. See Deploying models for additional instructions on deployment.

Note

A note on API compatibility

Since treelite is in early development, its API may change substantially in the future.

Benchmark

See the page Benchmark for details.