Treelite模型加速

本站内容均来自兴趣收集,如不慎侵害的您的相关权益,请留言告知,我们将尽快删除.谢谢.

今天有一朋友问我树模型算法预测的时候太慢怎幺办，我这里介绍一种基于树模型的模型加速工具，叫做Treelite。 Treelite : model compiler for decision tree ensembles

理论原理

该工作流涉及两台不同的机器：从给定树模型生成预测子例程的主机和运行该子例程的目标机器。这两台机器交换一个C文件，其中包含有关树模型的所有相关信息。只有主机需要安装Treelite；目标机器只需要一个正常工作的C编译器。

官方给出来的加速效果对比图，加速在2-7倍不等。

具体实践

定义数据集

import treelite
import numpy as np
import time
from sklearn.datasets import make_classification
X, y = make_classification(n_samples=10000, 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, 100, [(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)
    
    dmat = treelite_runtime.DMatrix(data)
    start = time.time()
    _ = predictor.predict(dmat)
    tree_lite.append(time.time() - start)
from matplotlib import pyplot as plt
plt.plot(nrows, xgb_time, "-*b", label="xgb")
plt.plot(nrows, tree_lite, "-*r", label="treelite")
plt.legend(loc="upper left")
plt.xticks(rotation=30)
plt.show()

模型训练轮数越多，treelite加速的程度越大。这里treelite模型加速了20倍。

LightGBM

LightGBM模型训练

import lightgbm as lgb
params = {
    'boosting_type': 'gbdt',
    'objective': 'binary',
    'max_depth':6,
    'metric': 'binary_logloss',
    'bagging_freq': 5,
    'verbose': 0
}
lgb_train = lgb.Dataset(X, y)
gbm = lgb.train(params,
                lgb_train,
                num_boost_round=100)
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)

TreeLite模型加载

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

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

nrows = [1000, 10000, 100000, 200000]
lgb_time = []
tree_lite = []
for nrow in nrows:
    data = np.random.random((nrow, 1000))
#     dtrain = lgb.Dataset(data)
    
    start = time.time()
    _ = gbm.predict(data)
    lgb_time.append(time.time() - start)
    
    dmat = treelite_runtime.DMatrix(data)
    start = time.time()
    _ = predictor.predict(dmat)
    tree_lite.append(time.time() - start)
plt.plot(nrows, lgb_time, "-*b", label="lgb")
plt.plot(nrows, tree_lite, "-*r", label="treelite")
plt.legend(loc="upper left")
plt.xticks(rotation=30)
plt.show()

模型训练轮数越多，treelite加速的程度越大。这里treelite模型加速了100倍。

Sklearn

随机森林模型训练

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

随机森林模型打包

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

TreeLite模型加载

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

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

nrows = [1000, 10000, 100000, 200000]
rf_time = []
tree_lite = []
for nrow in nrows:
    data = np.random.random((nrow, 1000))    
    start = time.time()
    _ = rf.predict(data)
    rf_time.append(time.time() - start)
    
    dmat = treelite_runtime.DMatrix(data)
    start = time.time()
    _ = predictor.predict(dmat)
    tree_lite.append(time.time() - start)
plt.plot(nrows, rf_time, "-*b", label="rf")
plt.plot(nrows, tree_lite, "-*r", label="treelite")
plt.legend(loc="upper left")
plt.xticks(rotation=30)
plt.show()

模型训练轮数越多，treelite加速的程度越大。这里treelite模型加速了50倍。

总结

1）文章重点展示了treelite方法的加速理论和实践，treelite模型可以大大加速预测，缩短响应代码；代码已经亲测，可以直接调用，如有问题，可随时留言。

2）重点参考了 阿水：Treelite：树模型部署加速工具（支持XGBoost、LightGBM和Sklearn） 。不过这个代码是一年前的，我复制下来运行的时候，会遇到小问题。在此向原作者阿水表示感谢。