How Treelite Supercharges Tree Model Inference by Up to 6×

What Is Treelite?

Treelite is an open‑source library that compiles decision‑tree models (XGBoost, LightGBM, scikit‑learn, etc.) into optimized shared libraries, enabling fast inference and easy deployment. Benchmarks show a 2‑6× speedup over native XGBoost predictions.

How Treelite Works

Treelite improves inference in two ways:

Branch Prediction

Tree traversal normally uses a series of if statements, causing the CPU to wait for each condition. By pre‑computing the number of samples that fall into each branch, Treelite can reorder branches and use compiler hints such as __builtin_expect to reduce mis‑predictions.

Comparison Simplification

Original implementations may compare floating‑point values. Treelite rewrites these as integer comparisons when possible, allowing the compiler to generate faster code.

Using Treelite

Below is a minimal workflow that works for any supported model.

Data Preparation

import treelite
import numpy as np
import time
from sklearn.datasets import make_classification

X, y = make_classification(n_samples=100000, n_features=1000)

XGBoost Example

Train a model

import xgboost
dtrain = xgboost.DMatrix(X, label=y)
params = {'max_depth':3, 'eta':1, 'objective':'reg:squarederror', 'eval_metric':'rmse'}
bst = xgboost.train(params, dtrain, 20, [(dtrain, 'train')])

Compile with Treelite

model = treelite.Model.from_xgboost(bst)
toolchain = 'gcc'   # adjust as needed
model.export_lib(toolchain=toolchain, libpath='./mymodel.so', verbose=True)

Load and predict

import treelite_runtime
predictor = treelite_runtime.Predictor('./mymodel.so', verbose=True)

Benchmark different batch sizes

nrows = [1000, 10000, 100000, 200000]
xgb_time = []
tree_lite = []
for nrow in nrows:
    data = np.random.random((nrow, 1000))
    dtrain = xgboost.DMatrix(data)

    start = time.time()
    _ = bst.predict(dtrain)
    xgb_time.append(time.time() - start)

    batch = treelite_runtime.Batch.from_npy2d(data)
    start = time.time()
    _ = predictor.predict(batch)
    tree_lite.append(time.time() - start)

LightGBM Example

Train a model

import lightgbm as lgb

params = {
    'boosting_type': 'gbdt',
    'objective': 'binary',
    'max_depth': 3,
    'metric': 'binary_logloss',
    'bagging_freq': 5,
    'verbose': 0
}

lgb_train = lgb.Dataset(X, y)
gbm = lgb.train(params, lgb_train, num_boost_round=10)

gbm.save_model('model.txt')

Compile with Treelite

model = treelite.Model.load('model.txt', model_format='lightgbm')
toolchain = 'gcc'   # adjust as needed
model.export_lib(toolchain=toolchain, libpath='./mymodel.so', verbose=True)

Scikit‑learn Example

Train a RandomForest model

from sklearn.ensemble import RandomForestClassifier
rf = RandomForestClassifier(n_estimators=1).fit(X, y)

Compile with Treelite

model = treelite.sklearn.import_model(rf)
toolchain = 'gcc'   # adjust as needed
model.export_lib(toolchain=toolchain, libpath='./mymodel.so', verbose=True)

Takeaways

Treelite provides a practical way to accelerate tree‑based inference and package models as shared objects.

It works with major frameworks (XGBoost, LightGBM, scikit‑learn) and also supports custom models.