Warpgbm Mcp

WarpGBM ⚡

神经网络般的速度。原生支持 GPU。感知分布。面向生产环境。

WarpGBM 是一款高性能、GPU 加速的梯度提升决策树（GBDT）库，基于 PyTorch 和自定义 CUDA 内核，从芯片层面精心打造。专为追求极致速度的用户和不愿妥协的研究人员而设计。

🎯 WarpGBM 的独特之处

统一支持回归与分类
使用同一套极速基础架构，训练连续目标或多类别标签。

不变性学习（DES 算法）
唯一能够原生学习在分布变化中保持稳定信号的开源 GBDT。由 Directional Era-Splitting 提供支持——因为你的数据并非存在于真空中。

全面 GPU 加速
为分箱、直方图、分裂和推理提供自定义 CUDA 内核。毫不妥协，不受 CPU 瓶颈限制。

兼容 Scikit-Learn
可直接替换。使用你熟悉的相同 API，获得你所需的 10 倍速度。

---

🚀 快速开始

安装

# Latest from GitHub (recommended)
pip install git+https://github.com/jefferythewind/warpgbm.git

# Stable from PyPI
pip install warpgbm

**先决条件：**支持 CUDA 的 PyTorch（安装指南）

5 行代码实现回归

from warpgbm import WarpGBM
import numpy as np

model = WarpGBM(objective='regression', max_depth=5, n_estimators=100)
model.fit(X_train, y_train)
predictions = model.predict(X_test)

5 行代码实现分类

from warpgbm import WarpGBM

model = WarpGBM(objective='multiclass', max_depth=5, n_estimators=50)
model.fit(X_train, y_train)  # y can be integers, strings, whatever
probabilities = model.predict_proba(X_test)
labels = model.predict(X_test)

---

🎮 功能

核心引擎

• ⚡ 原生 GPU CUDA 内核，用于构建直方图、查找分裂点、分箱和预测
• 🎯 支持多种目标：回归、二分类、多分类
• 📊 预分箱数据优化——如果数据已经量化，可跳过分箱
• 🔥 支持混合精度——输入类型为 float32 或 int8
• 🎲 随机化功能——使用 colsample_bytree 进行正则化

智能功能

• 🧠 通过 DES 实现不变性学习——识别能够跨时间、机制和环境泛化的信号
• 📈 智能初始化——分类使用类别先验，回归使用均值
• 🎯 自动标签编码——支持字符串、整数等各种输入
• 🔍 特征重要性——基于增益的重要性，并跟踪每个时期的独特重要性

训练工具

• ✅ 提前停止，支持验证集
• 📊 丰富的指标：MSE、RMSLE、相关性、对数损失、准确率
• 🔍 进度跟踪，显示损失曲线
• 🎚️ 正则化——叶节点 L2 惩罚、最小分裂增益、最小子节点权重
• 💾 热启动与检查点——保存/加载模型，增量训练

---

⚔️ 基准测试

合成数据：100 万行 × 1K 个特征（Google Colab L4 GPU）

   WarpGBM:   corr = 0.8882, train = 17.4s, infer = 3.2s  ⚡
   XGBoost:   corr = 0.8877, train = 33.2s, infer = 8.0s
  LightGBM:   corr = 0.8604, train = 29.8s, infer = 1.6s
  CatBoost:   corr = 0.8935, train = 392.1s, infer = 379.2s

速度是 XGBoost 的 2 倍，是 CatBoost 的 23 倍。

→ 亲自运行基准测试

多分类：3.5K 个样本、3 个类别、50 轮

Training:   2.13s
Inference:  0.37s
Accuracy:   75.3%

以神经网络般的速度实现面向生产环境的多分类。

---

📖 示例

回归：在你的笔记本电脑上击败 LightGBM

import numpy as np
from sklearn.datasets import make_regression
from warpgbm import WarpGBM

# Generate data
X, y = make_regression(n_samples=100_000, n_features=500, random_state=42)
X, y = X.astype(np.float32), y.astype(np.float32)

# Train
model = WarpGBM(
    objective='regression',
    max_depth=5, 
    n_estimators=100, 
    learning_rate=0.01,
    num_bins=32
)
model.fit(X, y)

# Predict
preds = model.predict(X)
print(f"Correlation: {np.corrcoef(preds, y)[0,1]:.4f}")

分类：带提前停止的多分类

from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split
from warpgbm import WarpGBM

# 5-class problem
X, y = make_classification(
    n_samples=10_000, 
    n_features=50,
    n_classes=5, 
    n_informative=30
)

X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2)

model = WarpGBM(
    objective='multiclass',
    max_depth=6,
    n_estimators=200,
    learning_rate=0.1,
    num_bins=32
)

model.fit(
    X_train, y_train,
    X_eval=X_val, y_eval=y_val,
    eval_every_n_trees=10,
    early_stopping_rounds=5,
    eval_metric='logloss'
)

# Get probabilities or class predictions
probs = model.predict_proba(X_val)  # shape: (n_samples, n_classes)
labels = model.predict(X_val)        # class labels

不变性学习：具备分布鲁棒性的信号

# Your data spans multiple time periods/regimes/environments
# Pass era_id to learn only signals that work across ALL eras

model = WarpGBM(
    objective='regression',
    max_depth=8,
    n_estimators=100
)

model.fit(
    X, y, 
    era_id=era_labels  # Array marking which era each sample belongs to
)

# Now your model ignores spurious correlations that don't generalize!

特征重要性：了解你的模型

from warpgbm import WarpGBM
from sklearn.datasets import load_iris

# Train a model
iris = load_iris()
X, y = iris.data, iris.target

model = WarpGBM(objective='multiclass', max_depth=5, n_estimators=100)
model.fit(X, y)

# Get feature importance (normalized)
importances = model.get_feature_importance()
for name, imp in zip(iris.feature_names, importances):
    print(f"{name}: {imp:.4f}")

# Output:
# sepal length (cm): 0.0002
# sepal width (cm): 0.0007
# petal length (cm): 0.1997
# petal width (cm): 0.7994

每个时期的特征重要性（WarpGBM 独有！）

使用 era_id 进行训练时，可以查看哪些特征在不同环境中保持稳定：

# Train with eras
model.fit(X, y, era_id=era_labels)

# Get per-era importance: shape (n_eras, n_features)
per_era_imp = model.get_per_era_feature_importance()

# Identify invariant features (high importance across ALL eras)
invariant_features = per_era_imp.min(axis=0) > threshold

热启动：增量训练与检查点

分阶段训练模型、保存检查点，并在之后恢复训练：

from warpgbm import WarpGBM
import numpy as np

# Train 50 trees
model = WarpGBM(
    objective='regression',
    n_estimators=50,
    max_depth=5,
    learning_rate=0.1,
    warm_start=True  # Enable incremental training
)
model.fit(X, y)
predictions_50 = model.predict(X_test)

# Save checkpoint
model.save_model('checkpoint_50.pkl')

# Continue training for 50 more trees (total: 100)
model.n_estimators = 100
model.fit(X, y)  # Adds 50 trees on top of existing 50
predictions_100 = model.predict(X_test)

# Or load and continue training later
model_loaded = WarpGBM()
model_loaded.load_model('checkpoint_50.pkl')
model_loaded.warm_start = True
model_loaded.n_estimators = 100
model_loaded.fit(X, y)  # Resumes from 50 → 100 trees

使用场景：

• 超参数调优：训练到 50 棵树后进行评估，再决定是否需要 100 棵或 200 棵
• 检查点：在长时间训练过程中保存进度
• 迭代开发：无需从头重新训练即可增加树的数量
• 生产环境更新：有新数据时增量重训练模型

预分箱数据：极速运行（Numerai 示例）

import pandas as pd
from numerapi import NumerAPI
from warpgbm import WarpGBM

# Download Numerai data (already quantized to integers)
napi = NumerAPI()
napi.download_dataset('v5.0/train.parquet', 'train.parquet')
train = pd.read_parquet('train.parquet')

features = [f for f in train.columns if 'feature' in f]
X = train[features].astype('int8').values
y = train['target'].values

# WarpGBM detects pre-binned data and skips binning
model = WarpGBM(max_depth=5, n_estimators=100, num_bins=20)
model.fit(X, y)  # Blazing fast!

结果：在 Numerai 数据上比 LightGBM 快 13 倍（49 秒 vs 643 秒）

---

🧠 不变性学习：为何重要

大多数机器学习模型都假设训练数据和测试数据来自相同的分布。现实情况是：并非如此。

• 股价会随市场状态变化
• 用户行为会随时间改变
• 实验数据会因批次／站点／条件而异

传统 GBDT： 会学习任何与目标相关的信号，包括在分布外（OOD）情况下会失效的脆弱模式。

使用 DES 的 WarpGBM： 会明确测试每个分裂能否在所有环境（时代）中泛化。只保留稳健的信号。

算法

对每个可能的分裂，分别计算其在每个时代中的增益。只接受满足以下条件的分裂：

1. 所有时代中的增益均为正
2. 各时代中的分裂方向一致

这样可以避免对只在某些时期或环境中有效的虚假相关性过拟合。

直观示意

左侧： 标准训练会将所有数据汇集在一起，学习任何相关信号。
右侧： 感知时代的训练要求信号在所有时期都有效，因此只学习稳健特征。

研究基础

• 不变风险最小化：Arjovsky 等，2019
• 难以变化的解释：Parascandolo 等，2020
• 树模型的时代划分：DeLise，2023

---

📚 API 参考

构造函数参数

WarpGBM(
    objective='regression',        # 'regression', 'binary', or 'multiclass'
    num_bins=10,                   # Histogram bins for feature quantization
    max_depth=3,                   # Maximum tree depth
    learning_rate=0.1,             # Shrinkage rate (aka eta)
    n_estimators=100,              # Number of boosting rounds
    min_child_weight=20,           # Min sum of instance weights in child node
    min_split_gain=0.0,            # Min loss reduction to split
    L2_reg=1e-6,                   # L2 leaf regularization
    colsample_bytree=1.0,          # Feature subsample ratio per tree
    random_state=None,             # Random seed for reproducibility
    warm_start=False,              # If True, continue training from existing trees
    threads_per_block=64,          # CUDA block size (tune for your GPU)
    rows_per_thread=4,             # Rows processed per thread
    device='cuda'                  # 'cuda' or 'cpu' (GPU strongly recommended)
)

训练方法

model.fit(
    X,                              # Features: np.array shape (n_samples, n_features)
    y,                              # Target: np.array shape (n_samples,)
    era_id=None,                    # Optional: era labels for invariant learning
    X_eval=None,                    # Optional: validation features
    y_eval=None,                    # Optional: validation targets  
    eval_every_n_trees=None,        # Eval frequency (in rounds)
    early_stopping_rounds=None,     # Stop if no improvement for N evals
    eval_metric='mse'               # 'mse', 'rmsle', 'corr', 'logloss', 'accuracy'
)

预测与实用方法

# Regression: returns predicted values
predictions = model.predict(X)

# Classification: returns class labels (decoded)
labels = model.predict(X)

# Classification: returns class probabilities
probabilities = model.predict_proba(X)  # shape: (n_samples, n_classes)

# Feature importance: gain-based (like LightGBM/XGBoost)
importances = model.get_feature_importance(normalize=True)  # sums to 1.0
raw_gains = model.get_feature_importance(normalize=False)   # raw gain values

# Per-era importance (when era_id was used in training)
per_era_imp = model.get_per_era_feature_importance(normalize=True)  # shape: (n_eras, n_features)

# Save and load models
model.save_model('checkpoint.pkl')  # Saves all model state
model_loaded = WarpGBM()
model_loaded.load_model('checkpoint.pkl')  # Restores everything

属性

model.classes_                    # Unique class labels (classification only)
model.num_classes                 # Number of classes (classification only)
model.forest                      # Trained tree structures
model.training_loss               # Training loss history
model.eval_loss                   # Validation loss history (if eval set provided)
model.feature_importance_         # Feature importance (sum across eras)
model.per_era_feature_importance_ # Per-era feature importance (when era_id used)

---

🔧 安装详情

Linux / macOS（推荐）

pip install git+https://github.com/jefferythewind/warpgbm.git

使用本地 PyTorch + CUDA 环境编译 CUDA 扩展。

Colab / CUDA 版本不匹配

pip install warpgbm --no-build-isolation

Windows

git clone https://github.com/jefferythewind/warpgbm.git
cd warpgbm
python setup.py bdist_wheel
pip install dist/warpgbm-*.whl

---

🎯 使用场景

金融机器学习： 学习在不同市场状态下均有效的信号
时间序列： 在分布变化下实现稳健预测
科学研究： 能够在不同实验批次间泛化的模型
高速推理： 满足毫秒级 SLA 的生产系统
Kaggle／竞赛： 使用 GPU 加速超参数调优
多分类问题： 图像分类后备方案、文本分类、欺诈检测

---

🚧 路线图

• 支持多 GPU 训练
• 在 GPU 上计算 SHAP 值
• 特征交互约束
• 单调性约束
• 自定义损失函数
• 分布式训练
• 导出 ONNX 以供部署

---

🙏 致谢

本项目构建于 PyTorch、scikit-learn、LightGBM、XGBoost 和 CUDA 生态系统之上。特别感谢 GBDT 研究社区和所有贡献者。

---

📝 版本历史

v2.2.0（当前）

• 💾 支持增量训练的热启动（解决 #14）
• 📦 用于检查点保存的 save_model() 和 load_model() 方法
• 🔄 使用 warm_start=True 从已保存的模型恢复训练
• ✅ 针对热启动和保存／加载功能的全面测试套件
• 📚 更新文档，加入热启动示例

v2.1.1

• 🎲 random_state 参数，用于确保结果可复现（解决 #12）
• 🔧 控制特征子采样（colsample_bytree）中的随机性
• ✅ 全面的可复现性测试

v2.1.0

• 🔍 特征重要性，支持基于增益的跟踪和独立的逐时代分析
• 📊 get_feature_importance() 和 get_per_era_feature_importance() 方法
• ✅ 与 LightGBM 对比的全面测试套件
• 📚 更新文档，加入特征重要性示例

v2.0.0

• ✨ 通过 softmax 目标函数支持多分类
• 🎯 二分类模式
• 📊 新增指标：对数损失、准确率
• 🏷️ 自动编码标签（支持字符串）
• 🔮 使用 predict_proba() 输出概率
• ✅ 全面的分类测试套件
• 🔒 与回归完全向后兼容
• 🐛 修复未使用变量问题（#8）
• 🧹 移除尚未实现的 L1_reg 参数
• 📚 全面更新文档，新增 AGENT_GUIDE.md

v1.0.0

• 🧠 通过方向性 Era 划分（DES）进行不变学习
• 🚀 VRAM 优化
• 📈 感知 Era 的直方图计算

v0.1.26

• 🐛 修复预测中的内存问题
• 📊 新增相关性评估指标

v0.1.25

• 🎲 特征子采样（colsample_bytree）

v0.1.23

• ⏹️ 支持提前停止
• ✅ 验证集评估

v0.1.21

• ⚡ CUDA 预测内核（替代向量化 Python 实现）

---

📄 许可证

MIT 许可证 - 请参阅 LICENSE 文件

---

🤝 参与贡献

欢迎提交 Pull Request！有关架构详情和开发指南，请参阅 AGENT_GUIDE.md。

---

由 @jefferythewind 🔥 打造

"训练更聪明。预测更快速。泛化更出色。"