Press "Enter" to skip to content

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) 。不过这个代码是一年前的,我复制下来运行的时候,会遇到小问题。在此向原作者阿水表示感谢。

Be First to Comment

发表评论

您的电子邮箱地址不会被公开。