hummingbird

项目介绍

什么是 Hummingbird？

Hummingbird 是一个将传统机器学习模型编译为张量计算的库。它允许用户无缝地利用神经网络框架（例如 PyTorch）来加速传统机器学习模型。这意味着用户可以通过 Hummingbird 来享受神经网络框架当前和未来所有的优化、硬件加速，并且在不需要重新设计模型的情况下，在一个统一的平台上支持传统和神经网络模型。

目前，Hummingbird 可以将训练好的传统机器学习模型转换为 PyTorch、TorchScript、ONNX 和 TVM。这些模型包括 scikit-learn 的决策树和随机森林，以及 LightGBM 和 XGBoost 分类器/回归器。Hummingbird 也计划在未来支持更多的神经网络后台和模型。

Hummingbird 提供了一个便捷且统一的“推理”API，遵循了 Sklearn 的 API。这意味着用户可以直接用 Hummingbird 生成的模型替换 Sklearn 模型，而无需更改推理代码。通过将模型转换为 PyTorch 和 TorchScript，用户还可以使用 TorchServe 提供服务。

Hummingbird 的工作原理

Hummingbird 的原理是将算法的操作符重新配置，以便进行向量化和 GPU 执行的常规计算。对于每一个操作符，其策略都稍有不同，但总体上主打的策略之一是将决策树翻译成张量涉及的 GEMM（通用矩阵乘法）。

以一个简单的决策树为例，该树有四个决策节点和五个叶节点。在输入一个具有五个元素的特征向量时，Hummingbird 通过矩阵乘法对树进行遍历，分三个步骤计算出最终的预测结果。在这一过程中，分别使用了四个张量（A、B、C、D、E）来表示输入特征、节点间关系、左右子树结构以及最终预测映射等。

安装指南

Hummingbird 已在 Python 3.9、3.10 和 3.11 的 Linux、Windows 和 macOS 系统上进行了测试（TVM 仅支持 Python 3.10）。建议使用虚拟环境来安装。Hummingbird 需要 PyTorch 版本大于等于 1.6.0。安装 PyTorch 后，可以使用以下命令通过 pip 安装 Hummingbird：

python -m pip install hummingbird-ml

若需要额外的 LightGBM 和 XGBoost 依赖，可以使用：

python -m pip install hummingbird-ml[extra]

使用示例

Hummingbird 的语法非常简洁直观。要在深度神经网络框架上运行传统机器学习模型，只需导入 hummingbird.ml 并在代码中添加 convert(model, 'dnn_framework') 模式。以下是一个使用 scikit-learn 随机森林模型和 PyTorch 作为目标框架的示例：

import numpy as np
from sklearn.ensemble import RandomForestClassifier
from hummingbird.ml import convert, load

# 创建随机数据用于二分类
num_classes = 2
X = np.random.rand(100000, 28)
y = np.random.randint(num_classes, size=100000)

# 创建并训练一个模型（这里使用的是 scikit-learn 随机森林分类器）
skl_model = RandomForestClassifier(n_estimators=10, max_depth=10)
skl_model.fit(X, y)

# 使用 Hummingbird 将模型转换为 PyTorch
model = convert(skl_model, 'pytorch')

# 在 CPU 上运行预测
model.predict(X)

# 在 GPU 上运行预测
model.to('cuda')
model.predict(X)

# 保存模型
model.save('hb_model')

# 重新加载模型
model = load('hb_model')