Importing tree ensemble models

Since the scope of Treelite is limited to prediction only, one must use other machine learning packages to train decision tree ensemble models. In this document, we will show how to import an ensemble model that had been trained elsewhere.

Importing XGBoost models

XGBoost (dmlc/xgboost) is a fast, scalable package for gradient boosting. Both Treelite and XGBoost are hosted by the DMLC (Distributed Machine Learning Community) group.

Treelite plays well with XGBoost — if you used XGBoost to train your ensemble model, you need only one line of code to import it. Depending on where your model is located, you should do one of the following:

# bst = an object of type xgboost.Booster
model = Model.from_xgboost(bst)
  • Load XGBoost model from a binary model file

# model had been saved to a file named my_model.model
# notice the second argument model_format='xgboost'
model = Model.load('my_model.model', model_format='xgboost')

Importing LightGBM models

LightGBM (Microsoft/LightGBM) is another well known machine learning package for gradient boosting. To import models generated by LightGBM, use the load() method with argument model_format='lightgbm':

# model had been saved to a file named my_model.txt
# notice the second argument model_format='lightgbm'
model = Model.load('my_model.txt', model_format='lightgbm')

Importing scikit-learn models

Scikit-learn (scikit-learn/scikit-learn) is a Python machine learning package known for its versatility and ease of use. It supports a wide variety of models and algorithms. The following kinds of models can be imported into Treelite.

To import scikit-learn models, use treelite.sklearn.import_model():

# clf is the model object generated by scikit-learn
import treelite.sklearn
model = treelite.sklearn.import_model(clf)

How about other packages?

If you used other packages to train your ensemble model, you’d need to specify the model programmatically. There are two ways to do this: