模型部署（5）Treelite加速

发表于 更新于 分类于 ， 阅读次数： 本文字数： 1.1k 阅读时长 ≈ 2 分钟

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

一、TreeLite介绍

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

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

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

二、TreeLite原理

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

  • 逻辑分支

    • 这个指令是gcc引入的，作用是允许程序员将最有可能执行的分支告诉编译器。这个指令的写法为：__builtin_expect(EXP, N)。意思是：EXP==N的概率很大。【减少重新取跳转地址】__builtin_expect 说明: https://www.jianshu.com/p/2684613a300f

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

  • 逻辑比较

    • 浮点数比较，转化为整数型比较。

2.1 逻辑分支

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

1
2
3
4
5
if ( [conditional expression] ) {
  foo();
} else {
  bar();
}

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

img

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

1
2
3
4
5
6
/* expected to be false */
if( __builtin_expect([condition],0)){
  ...
} else {
  ...
}

2.2 逻辑比较

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

1
2
3
4
5
6
7
8
9
if (data[3].fvalue < 1.5) {  
/* floating-point comparison */
  ...
}

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

参考文献