Treelite：树模型部署加速工具（支持XGBoost、LightGBM和Sklearn）

项目链接：https://treelite.readthedocs.io/

项目论文：https://mlsys.org/Conferences/doc/2018/196.pdf

支持模型：XGB、LGB、SKlearn树模型

Treelite介绍

Treelite能够树模型编译优化为单独库，可以很方便的用于模型部署。经过优化后可以将XGBoost模型的预测速度提高2-6倍。

如上图，黑色曲线为XGBoost在不同batch size下的吞吐量，红色曲线为XGBoost经过TreeLite编译后的吞吐量。

Treelite支持众多的树模型，特别是随机森林和GBDT。同时Treelite可以很好的支持XGBoost, LightGBM和 scikit-learn，也可以将自定义模型根据要求完成编译。

Treelite原理

Treelite主要在两方面完成了改进。

逻辑分支

对于树模型而言，节点的分类本质使用if语句完成，而CPU在执行if语句时会等待条件逻辑的计算。

if ( [conditional expression] ) {
  foo();
} else {
  bar();
}

如果在构建树模型时候，提前计算好每个分支下面样本的个数，则可以提前预知哪一个叶子节点被执行的可能性更大，进而可以提前执行子节点逻辑。

借助于编译命令，可以完成逻辑计算加速。

/* expected to be false */
if( __builtin_expect([condition],0)){
  ...
} else {
  ...
}

逻辑比较

原始的分支比较可能会有浮点数比较逻辑，可以量化为数值比较逻辑。

if (data[3].fvalue < 1.5) {  
/* floating-point comparison */
  ...
}

if (data[3].qvalue < 3) {     
/* integer comparison */
  ...
}

Treelite使用

首先定义数据集：

import treelite
import numpy as np
import time
from sklearn.datasets import make_classification
X, y = make_classification(n_samples=100000, n_features=1000)

Treelite模型打包后，加载和之后的预测代码一样，也就说所有树模型在TreeLite打包后之后的语法都是一样的。

XGBoost

• XGBoost模型训练

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')])

• XGBoost模型打包

model = treelite.Model.from_xgboost(bst)
toolchain = 'gcc'   # change this value as necessary
model.export_lib(toolchain=toolchain, libpath='./mymodel.so', verbose=True)

• Treelite模型加载

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

• XGBoost不同Batch大小预测速度对比

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

• LightGBM模型训练

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')

• LightGBM模型打包

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

• LightGBM不同Batch大小预测速度对比

Sklearn

• 随机森林模型训练

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

• 随机森林模型打包

# model = treelite.Model.from_xgboost(bst)
model = treelite.sklearn.import_model(rf)
toolchain = 'gcc'   # change this value as necessary
model.export_lib(toolchain=toolchain, libpath='./mymodel.so', verbose=True)

• 随机森林不同Batch大小预测速度对比

后记

1. Treelite是不错的树模型加速部署工具；
2. Treelite支持将模型部署为.so文件，方便其他语言调用，也可以避免模型参数泄露；