lleaves：使用 LLVM 编译梯度提升决策树将预测速度提升10+倍

梯度提升决策树是一种常用的机器学习算法，在真实世界的表格数据集上表现良好。有许多库可用于训练它们，最常见的是 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 编译这个模型很简单：

ini 复制代码

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 文件中，该文件看起来像这样：

ini 复制代码

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：

vbnet 复制代码

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 上，树的函数现在看起来像这样（被截断）：

ini 复制代码

.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中相关代码如下所示：

arduino 复制代码

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。

原文链接：siboehm.com/articles/21...