梯度提升决策树是一种常用的机器学习算法,在真实世界的表格数据集上表现良好。有许多库可用于训练它们,最常见的是 LightGBM 和 XGBoost。 遗憾的是,很少有流行的库针对快速预测和部署进行优化。 作为补救措施,我花了几个月的时间构建了 lleaves,一个开源决策树编译器和 Python 包。
它可以用作 LightGBM 的直接替代品,并将推理速度提高 ≥10 倍。 lleaves 通过 LLVM 编译决策树以生成优化的机器代码。 目标是让 lleaves 为每个数据科学家提供最先进的预测加速。
环境准备
查看GitHub仓库:lleaves,并通过 pip install lleaves
或 conda install -c conda-forge lleaves
安装即可!
为什么要使用lleaves?
lleaves 速度快、易于使用,并且支持 LightGBM 的所有标准功能。 下面我将使用在 MTPL2 数据集上训练的 LightGBM 模型进行演示。 使用 lleaves 编译这个模型很简单:
model = lleaves.Model(model_file="MTPL2/model.txt")
model.compile()
result = model.predict(df)
运行时间从 9.7 秒 (LightGBM) 下降到 0.4 秒 (leleaves)。如果您交互地使用树模型,这种加速会产生相当明显的差异。
这里我们使用多线程对大batches进行预测,这些基准测试将 lleaves 与 ONNX 和 Treelite 进行比较,这两个库编译决策树以在 CPU 上进行预测,具体如下图所示。
当然,lleaves 同样适用于小批量的单线程预测。
lleaves 是如何工作的?
lleaves 是一个编译器,这意味着它将代码从一种语言转换为另一种语言。 更准确地说:lleaves 是 LLVM 编译器工具包的前端,将 LightGBM 的 model.txt 转换为 LLVM IR。 最后,我们需要将训练好的模型转化为 CPU 可以理解的汇编代码。 完整的编译分 3 个步骤进行:
- LightGBM 将训练后的模型存储在磁盘上的 model.txt 文件中。
- lleaves 加载 model.txt 并将其转换为 LLVM IR。
- LLVM 将 LLVM IR 转换为本地汇编。
lleaves 严重依赖 LLVM 来生成汇编(assembly)。 LLVM 是一个“编译器工具包”,这意味着它提供了模块和部件来使构建编译器变得更容易。 其中一个部分是 LLVM 中间表示 (IR),这是一种类似于 RISC 汇编的低级编程语言。 IR还不是机器代码,因此仍然需要进行转换才能被CPU执行。
lleaves 的架构与许多现代编译器类似,如:clang (C/C++)、rustc (Rust) 或 numba (Python),但复杂程度要低得多。
现代编译器架构
现代编译器由三个阶段组成,如下所示:
- 前端将源代码转换为独立于硬件架构的中间表示(IR)。
- 优化器获取生成的 IR 并对其运行优化过程。这些 pass 再次输出 IR,但可能会删除冗余代码或用相同但更快的操作替换昂贵的操作。
- 优化后的 IR 被传递到编译器的后端,将 IR 转换为汇编语言。 在此步骤中还会进行特定于硬件架构的优化。
LLVM 工具包规定了 LLVM IR,并包括优化器和后端。 这意味着对于 LLVM,编写新编译器所需的只是提供一个将源代码转换为 LLVM IR 的前端。 然后LLVM的优化器和后端调整IR并将其转换为汇编。
实际上,它稍微复杂一些。 例如,许多编译器(如:numba)包含自己的、更高级的 IR,以便在最终降低到 LLVM IR 之前更轻松地运行特定于语言的优化。
还记得我说过 lleaves 是 LLVM 的前端吗? 这是因为 lleaves 将 LightGBM 输出的 model.txt 转换为 LLVM IR。 然后 lleaves 将 IR 传递给 LLVM 的优化器和后端,后者处理机器代码(汇编)的最终转换。
步骤 1:LightGBM –> Model.txt
为了说明这一点,我们将遵循经过训练的 LightGBM 模型完成从 model.txt 到汇编(assembly)的每个转换步骤。 作为一个运行示例,接下来看一下使用 LightGBM 训练的单个决策树:
树中的每个节点将输入值与其阈值进行比较。 如果该值小于或等于阈值,则采用左侧路径。 否则,采用右侧路径。 一旦我们到达叶子,叶子的值就会被返回。
LightGBM 将我们的示例树存储在 model.txt 文件中,该文件看起来像这样:
Tree=0
num_leaves=4
threshold=0.731 0.907 0.856 # each node's threshold
split_feature=1 2 2 # input feature to compare threshold with
left_child=1 -1 -2 # for each node, index of left child
right_child=2 -3 -4 # for each node, index of right child
leaf_value=0.495 0.507 0.506 0.490 # each leaf's return value
对于任何现实世界的模型,这里显然会有 100 多棵树,而不仅仅是一棵。
步骤 2:Model.txt –> LLVM IR
lleaves 从磁盘加载此 model.txt,解析它,并将树转换为 LLVM IR:
define private double @tree_0(double %.1, double %.2, double %.3) {
node_0:
%.5 = fcmp ule double %.2, 0x3FE768089A419B12 ; decimal = ~0.731
br i1 %.5, label %node_1, label %node_2
node_1: ; preds = %node_0
%.7 = fcmp ule double %.3, 0x3FED06D4513F4FE5 ; decimal = ~0.907
br i1 %.7, label %leaf_0, label %leaf_2
node_2: ; preds = %node_0
%.11 = fcmp ule double %.3, 0x3FEB60631F166F7A ; decimal = ~0.856
br i1 %.11, label %leaf_1, label %leaf_3
leaf_0: ; preds = %node_1
ret double 0x3FDFAFD3A55B8741 ; decimal = ~0.495
leaf_2: ; preds = %node_1
ret double 0x3FE038704B651588 ; decimal = ~0.507
leaf_1: ; preds = %node_2
ret double 0x3FE034DEA54DFC96 ; decimal = ~0.506
leaf_3: ; preds = %node_2
ret double 0x3FDF62CFF241EA8B ; decimal = ~0.490
}
在 IR 片段中,您可以看到树的每个节点是如何处理的。 输入数据通过函数的属性(%.1、%.2、%.3)传递。 相关属性与阈值进行浮点比较(fcmp),以检查该属性是否 unsigned less-or-equal(ule)阈值。每个双精度数都已转换为其十六进制表示形式。 根据结果,我们跳转/中断 (br) 到下一个节点,最后返回 (ret) 结果。
该模型现已转换为低级标准化语言。 剩下要做的就是将其编译为汇编。
步骤 3:LLVM IR –> 汇编
在 lleaves 将模型转换为 LLVM IR 后,LLVM 对其运行优化 pass 并调整本机 CPU 架构的代码。 然后生成最终的汇编代码。 在 Intel CPU 上,树的函数现在看起来像这样(被截断):
.LCPI0_0:
.quad 0x3fe768089a419b12 ; double 0.731
.LCPI0_1:
.quad 0x3feb60631f166f7a ; double 0.856
.LCPI0_3:
.quad 0x3fed06d4513f4fe5 ; double 0.907
.LCPI0_2:
.quad 0x3fdf62cff241ea8b ; double 0.490
.quad 0x3fe034dea54dfc96 ; double 0.506
.LCPI0_4:
.quad 0x3fe038704b651588 ; double 0.507
.quad 0x3fdfafd3a55b8741 ; double 0.495
tree_0: ; @tree_0
xor eax, eax
ucomisd xmm1, qword ptr [rip + .LCPI0_0]
ja .LBB0_2
ucomisd xmm2, qword ptr [rip + .LCPI0_3]
setbe al
movsd xmm0, qword ptr [8*rax + .LCPI0_4]
ret
.LBB0_2: ; %node_2
ucomisd xmm2, qword ptr [rip + .LCPI0_1]
setbe al
movsd xmm0, qword ptr [8*rax + .LCPI0_2]
ret
如果您不习惯的话,汇编很难阅读。 将此与上面的 LLVM IR 进行比较,后者更容易理解(因此也更容易生成!)
总结三个编译步骤: 使用 LightGBM 训练模型并保存到 model.txt。 然后 lleaves 加载 model.txt 并将模型转换为与硬件架构无关的 LLVM IR。 IR 被传递到 LLVM,针对特定 CPU 架构进行优化,并转换为汇编。 然后可以从 Python 调用最终的程序集,传入数据并返回结果。
为什么 lleaves 的预测速度比 LightGBM 更快?
在LightGBM中相关代码如下所示:
inline int NumericalDecision(double fval, int node) const {
if (fval <= threshold_[node]) {
return left_child_[node];
}
return right_child_[node];
}
LightGBM 不会为其训练模型生成任何代码。 相反,model.txt 在运行时从磁盘加载到内存中。 然后,通过对树中的每个节点调用一次 NumericalDecision,逐个节点执行树。 每棵树和每个节点的决策函数都保持完全相同。 lleaves 内部的模型编译工作方式非常不同,它为每个决策节点生成唯一的代码。 此编译步骤允许编译时优化、消除内存访问开销并解锁现代 CPU 功能,例如:分支预测。
结论
在小于 100 棵树到大于 1000 棵树的各种模型上测试 lleaves,加速速度从 10 倍到 30 倍不等。 如果您在部署中使用 LightGBM,可以尝试一下 lleaves。