让树模型预测时能加速2-10倍的方案
↑↑↑关注后"星标"kaggle竞赛宝典
kaggle竞赛宝典 作者：Coggle
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大小预测速度对比
后记
- Treelite是不错的树模型加速部署工具；
- Treelite支持将模型部署为.so文件，方便其他语言调用，也可以避免模型参数泄露；