English | 中文
leaves 是一个纯 Go 写的推荐系统用排序模型训练与推理库。能直接加载 XGBoost / LightGBM / scikit-learn 模型,也能自训练(hist / exact、排序、生存、tweedie 等),全程不需要任何 C 绑定。
推荐入口
- 推理:
leaves.LoadFromFile(默认开启AutoTransform),或者沿用旧的leaves.LGEnsembleFromFile/leaves.XGEnsembleFromFile。 - 训练:
leaves.NewLearner/train.NewLearner、leaves.LoadDataAuto/data.FromFileAuto(带内容嗅探),以及便利函数leaves.NewLearnerFromModelAndData—— 它能从参考模型里反推出 objective。
回归矩阵见 docs/testdata-matrix.md;路线图 演进计划.md(v5.4);Agent 契约 演进方案.md(v1.5);扩展 objective/metric 见 docs/extension-points.md;backlog TODO.md。
- 并行批预测(dense / CSR)
- 支持 sigmoid、softmax 变换;
DefaultLoadOptions()默认开启 AutoTransform - 可输出决策树的叶节点索引
- 训练数据内容嗅探(LIBSVM、排序 TSV、CSV·TSV —— 扩展名只是回退依据)
- 兼容
text与JSON两种模型格式 - 支持
gbdt、rf(随机森林)、dart三类模型 - 支持多分类预测
- 对类别特征(独热分割等)有专门的决策规则优化
- 纯预测场景下的快速路径
- 支持二进制、JSON、UBJSON 三种模型格式
gbtree/gblinear/dart三种 booster;3.x 嵌套 DARTgbtree.model- 多分类预测、多输出向量叶、XGBoost 原生 categorical(
model.cats)并支持 round-trip 导出 - 缺失值(
nan)支持 reg:gamma/count:poisson/reg:tweedie自动加载变换(Exp +base_score对数转换)
scikit-learn(实验性)
- 支持
pickle协议0 - 兼容
sklearn.ensemble.GradientBoostingClassifier
gbtree/dart/gblinear;hist/exact/gpu_hist(auto模式按行数自动选)- 目标函数:回归、二分类 / 多分类、gamma、poisson、tweedie、
survival:cox/survival:aft,以及排序(rank:ndcg/rank:pairwise/rank:listwise) - 交叉验证、早停、checkpoint、续训、单调约束、
max_leaves(lossguide)、BatchedMatrix(外存 DMatrix) - 模型导出:原生
leaves.json、XGBoost 3.x JSON
- int8 阈值量化(带 parity 门禁)
Ensemble.Reload线上热更新;按 batch 维度的推理 profiling- WASM demo、HTTP embed、serving 模板、train-from-model、扩展 objective/metric、多目标回归
- MovieLens Ranker + Agent/MCP 教程(精排全流程案例)
go install github.com/linkerlin/leaves@latest模块路径:github.com/linkerlin/leaves(Go 1.26+)。张量与 GPU 加速由 github.com/born-ml/born 提供。
最简单的方式:用 leaves.LoadFromFile 加载模型,然后调 PredictSingle。
package main
import (
"fmt"
"github.com/linkerlin/leaves"
)
func main() {
m, err := leaves.LoadFromFile("lightgbm_model.txt", leaves.DefaultLoadOptions())
if err != nil {
panic(err)
}
fvals := []float64{1.0, 2.0, 3.0}
p, err := m.PredictSingle(fvals, 0)
if err != nil {
panic(err)
}
fmt.Printf("对 %v 的预测: %f\n", fvals, p)
}旧 API 仍兼容:leaves.LGEnsembleFromFile(path, loadTransformation) 和 leaves.XGEnsembleFromFile(...),行为不变。
io.LoadFromFile / DetectFormat 看扩展名与文件头。完整等级表见 docs/interop-matrix.md(稳定 / 实验 / 占位)。
| 格式 | 等级 | 扩展名 / 探测 | 说明 |
|---|---|---|---|
| leaves.json | 稳定 | leaves_version |
训练主产物;推荐部署 |
| XGBoost JSON | 稳定 | .json + learner |
3.x 默认;base_score→margin |
| XGBoost UBJSON | 稳定 | .ubj |
与 JSON 预测一致 |
| XGBoost binary | 稳定 | binf / header |
经典 Booster;优先改用 JSON |
| LightGBM text/JSON | 稳定 | tree= / tree_info |
text 与 JSON |
| scikit-learn | 实验 | .pkl / .joblib |
窄协议;生产请转 XGB/leaves JSON |
| ONNX | 实验 | .onnx |
TreeEnsembleRegressor 子集;复杂图请先转 JSON/leaves |
加载失败返回 *io.LoadError(含 hint: 下一步)。数值表误用 .txt 会提示改用 data.FromFile。
import "github.com/linkerlin/leaves/io"
m, err := io.LoadFromFile("model.ubj", io.DefaultLoadOptions()) // AutoTransform = true
m, err = io.LoadFromFile("model.ubj", &io.LoadOptions{
AutoTransform: false, // 输出原始 margin
LoadTransformation: true, // 或者对所有目标强制启用变换
Backend: io.BackendAuto,
})当 AutoTransform 与 LoadTransformation 都为 false 时,OutputValue 和 OutputMargin 都是原始 margin。contrib / SHAP 始终在 margin 空间,与 AutoTransform 无关。
dm, _ := leaves.LoadDataAuto("train.tsv") // 内容嗅探
learner, _ := leaves.NewLearner(leaves.TrainConfig{
Objective: leaves.TrainObjectiveBinaryLogistic,
NumRound: 50,
})
_ = learner.Fit(dm)
_ = learner.Save("out.leaves.json")bm, _ := data.SplitDense(dense, batchRows) // 也可以 data.NewBatchedMatrix(...)
learner, _ := train.FitExternal(train.Config{
Objective: train.ObjectiveBinaryLogistic,
TreeMethod: train.TreeMethodHist,
HistBinPolicy: "global",
NumRound: 10,
}, bm)learner, _ := train.ResumeFit("ckpt.json", train.Config{
Objective: train.ObjectiveSquaredError,
NumRound: 20, // 在 checkpoint 完成的轮次之上继续
TreeMethod: train.TreeMethodExact,
}, dm)
score, _ := learner.Eval(dm)data.FromFileAuto / leaves.LoadDataAuto 会先读一段样本来选 loader,扩展名只是兜底。
| 嗅探结果 | 典型内容 | Loader |
|---|---|---|
| LIBSVM | label 1:0.5 3:1.2 |
FromLIBSVM |
| 排序 TSV | qid label feat... |
LoadRankingTSV |
| 末列为 label 的 TSV | 无表头,末列是 label | LoadDenseTSV |
| CSV | 有表头或启发式 label 列 | FromCSV |
learner, _ := leaves.NewLearnerFromModelAndData(
"reference_model.json", // 只读 objective(XGB JSON / UBJ / leaves.json)
"train.tsv", // LoadDataAuto 嗅探
leaves.TrainConfig{NumRound: 10, TreeMethod: leaves.TrainTreeMethodExact},
data.DefaultFileLoadOptions(),
)
_ = learner.Save("out.leaves.json")go run ./examples/train_from_model/ -model testdata/xgboost_smoke.json -data testdata/breast_cancer_train.tsvd, _ := data.NewDense(vals, n, nfeat, placeholderLabels, nil)
dm, _ := data.NewAFTDense(d, []data.AFTInterval{
{1.5, 1.5}, // 事件
{2, math.Inf(1)}, // 右删失
{0, 4}, // 左删失
{1, 6}, // 区间删失
})
learner, _ := train.NewLearner(train.Config{Objective: train.ObjectiveSurvivalAFT, ...})
_ = learner.Fit(dm)learner, _ := train.NewLearner(train.Config{
Objective: "reg:squarederror",
MonotoneConstraints: []int{1, 0, -1}, // 特征 0 递增,特征 2 递减
TreeMethod: train.TreeMethodHist,
})hist / gpu_hist 路径上的加速分层(日志前缀 [leaves/train] accel:):
| 阶段 | 实现 | 默认 auto(≥3 万行 + WebGPU) |
LEAVES_TRAIN_ACCEL=webgpu |
|---|---|---|---|
| 全局分箱 | 训练期每特征一次切点 + 行级 bin 缓存 | ✅ | ✅ |
| 直方图累加 | CPU rowBin / WebGPU SelectAdd 批量 |
gpu_hist |
GPU batch |
| 增益扫描 | 纯 CPU / 批量 WebGPU gain | 批量 GPU(GPU 路径) | 批量 GPU |
| margin 预测 | 逐行 CPU / Born GPU PredictDense |
PredictMargins ≥256 行 |
≥64 行 |
learner, _ := train.NewLearner(train.Config{
Objective: train.ObjectiveRankNDCG,
TreeMethod: train.TreeMethodGPUHist,
AccelMode: train.AccelModeWebGPU, // 也可走环境变量 LEAVES_TRAIN_ACCEL
NumThreads: 4,
HistBinPolicy: "global", // 默认;"per_node" 恢复旧的 per-node 切点
NumRound: 50,
})
_ = learner.Fit(dm)环境变量覆盖:LEAVES_TRAIN_ACCEL=auto|webgpu|born_cpu|cpu。
计算底座 —— Born
推理与训练加速都构建在 Born 之上(CPU SIMD + WebGPU)。NativeEngine 是 golden 基准;BackendAuto(2.0)按 workload 派发 BornCPU / BornGPU。完整决策表:docs/backend-auto.md。
m, _ := leaves.LoadFromFile("model.json", &io.LoadOptions{
Backend: io.BackendAuto,
Workload: tree.WorkloadHint{BatchSize: 256, HasGPU: true},
})
// 可解释选型:tree.SelectBackendExplained(caps, hint) → Rule / Reasonparity 门禁 TestBornParityFormatMatrix 覆盖 LGB text/JSON、XGB bin/json/ubj、scikit-learn pickle × batch {1, 16, 256} × BornCPU / BornGPU,容差 1e-5(相对 Native)。
| 场景 | Auto 结果 | 说明 |
|---|---|---|
| 在线单条 / batch < 64 | Native | 延迟优先;Born 单条更慢 |
| 批量 ≥ 64,纯数值树 | BornCPU | SIMD 批推理 |
| 批量 ≥ 256 + HasGPU(Windows WebGPU) | BornGPU | 不可用则回落 BornCPU |
| WASM + 数值树 | BornCPU(支持时) | 无 GPU;cat-small → Native |
稀疏 SparseDensity∈(0,0.15) / cat-small |
Native | Born 未优化 / 不支持 |
m.Explain() 暴露的是 Lundberg 快速 Tree SHAP(带 SumHess 权重覆盖,tree_path_dependent);SHAP 值在 margin 空间,与 OutputMargin / predict.OutputContribution 保持一致。DefaultLoadOptions() 对 logistic 模型返回概率;contrib / SHAP 仍在 margin 空间。
import (
"github.com/linkerlin/leaves"
"github.com/linkerlin/leaves/explain"
)
m, _ := leaves.LoadFromFile("model.json", leaves.DefaultLoadOptions())
x := [][]float64{{1.0, 2.0, 3.0}}
contrib, _ := m.Explain().TreeSHAP(x) // margin 空间 SHAP
inter, _ := m.Explain().InteractionSHAP(x) // 交互值矩阵
mc, _ := m.Explain().TreeSHAPMulticlass(x) // 多类:[sample][feature][class]
base := m.Explain().ExpectedValue() // 背景基线(全零特征)
imp := m.Explain().Importance(explain.ImportanceGain, nil)
text := m.Explain().DumpText(nil)
dot := m.Explain().DumpDOT(nil)统一出口是 predict.Request:
import "github.com/linkerlin/leaves/predict"
nf := m.NFeatures()
nRows := 1
// 输出值(logistic 已变换,squarederror 仍为原始)
vals := make([]float64, nRows*m.NOutputGroups())
_ = m.PredictWithRequest(predict.Request{
Matrix: predict.DenseMatrix{Values: flat, Rows: nRows, Cols: nf},
Output: predict.OutputValue,
}, vals)
// 原始 margin(sigmoid / softmax 之前)
margins := make([]float64, nRows*m.NRawOutputGroups())
_ = m.PredictWithRequest(predict.Request{
Matrix: predict.DenseMatrix{Values: flat, Rows: nRows, Cols: nf},
Output: predict.OutputMargin,
}, margins)
// 叶节点索引(每棵树一个 int32)
nTrees := m.NTrees()
leaves := make([]int32, nRows*nTrees)
_ = m.PredictWithRequest(predict.Request{
Matrix: predict.DenseMatrix{Values: flat, Rows: nRows, Cols: nf},
Output: predict.OutputLeaf,
}, leaves)
// SHAP 贡献(margin 空间;二分类 [sample][feature+bias])
out := make([]float64, (nf+1)*nRows)
_ = m.PredictWithRequest(predict.Request{
Matrix: predict.DenseMatrix{Values: flat, Rows: nRows, Cols: nf},
Output: predict.OutputContribution, // 也可 OutputApproxContribution / OutputInteraction
}, out)bias 语义:末列是全零特征处的背景 margin,与 XGBoost pred_contribs 的可加性约定一致(逐元素分解可能不同)。
内置 RMSE / MAE / AUC / LogLoss / MAPE / RMSLE / NDCG@k / MAP,命名与 XGBoost eval_metric 对齐:
import "github.com/linkerlin/leaves/metrics"
rmse, _ := metrics.RMSE{}.Evaluate(yTrue, yPred)
m, _ := metrics.Resolve("ndcg@5", metrics.Options{Groups: []int{10, 10}})
ndcg, _ := m.Evaluate(yTrue, yPred)排序指标需要 data.GroupedMatrix 提供 Groups()。
XGBoost 兼容的 LambdaMART,再加上原生的 listwise 头:
| 目标 | 家族 | 说明 |
|---|---|---|
rank:ndcg |
LambdaMART(XGBoost 兼容) | pairwise + |ΔNDCG| 缩放,直接对齐排序指标 |
rank:pairwise |
RankNet pairwise | 默认 top-k(k=32)+ lambdarank_normalization(对标 XGBoost) |
rank:listwise |
ListNet softmax CE(leaves 原生) | q ∝ exp(label)、p = softmax(pred),纯 listwise 交叉熵 |
import (
"github.com/linkerlin/leaves/data"
"github.com/linkerlin/leaves/train"
)
dm, _ := data.LoadRankingTSV("rank_train.tsv", "\t") // qid label feat1 feat2 ...
learner, _ := train.NewLearner(train.Config{
Objective: train.ObjectiveRankNDCG, // 或 ObjectiveRankPairwise / ObjectiveRankListwise
NumRound: 40,
MaxDepth: 4,
LearningRate: 0.1,
NDCGK: 10,
EvalMetric: "ndcg@10",
})
_ = learner.Fit(dm)端到端 MovieLens 100K demo(训练 → 保存 → Top-K 推荐):demos/movielens/README.md。
cd testdata && python gen_rank_movielens.py
go run ./demos/movielens/cmd/train
go run ./demos/movielens/cmd/recommend -group 0 -topk 10
go test ./train/... -run TestRankMovieLens -vlearner, _ := train.NewLearner(train.Config{
Objective: train.ObjectiveSquaredError,
NumRound: 20,
LearningRate: 0.3,
LRScheduler: train.ExponentialLRScheduler(0.95), // 每轮 ×0.95
Callbacks: []train.TrainingCallback{
train.FuncCallback(func(ctx *train.CallbackContext) error {
// ctx.Round, ctx.LearningRate, ctx.TrainMetric, ctx.EvalMetric
return nil
}),
},
})内置调度器:ExponentialLRScheduler(gamma) 和 StepLRScheduler(every, factor)。配置 EvalSet + EvalMetric 时,CallbackContext 也会自动填充 EvalMetric / EvalMetricOK(不依赖 EarlyStop)。
if ne, ok := m.Engine().(*tree.NativeEngine); ok {
prof, err := tree.ProfileNativeDense(ne, vals, nrows, ncols, preds, 0)
_ = prof.Elapsed
}
// 线上原子切换模型
_ = m.Reload("/path/to/new.model", io.DefaultLoadOptions())tree.ProfileWalkStats 可以单独统计单样本的树遍历深度,不跑完整预测。
数值分裂阈值按特征做 int8 仿射量化(127 级),叶子值保持 float64,分类节点不量化。quantize.Engine 不支持 PredictLeafIndices*(只支持 margin 预测)。
qf, _ := quantize.QuantizeForest(m.Forest(), quantize.Config{})
res, err := quantize.CheckParityWithGate(m.Forest(), qf, samples, 0, quantize.DefaultGate())
eng, _ := quantize.NewEngine(qf, nil, tree.TransformRaw, m.NOutputGroups())
model.NewEnsemble(eng) // 替换线上 Ensemble 引擎Agentic 契约已达成(见
演进方案.mdv1.5 DoD)。
库整体 12 个月路线图见演进计划.mdv5.4 — 二者互补,不互相替代。
Agent 通过 SKILL + leaves CLI + metrics.json 完成「训练→调参→发布」;Agent 即优化器(搜索逻辑在 SKILL,库不内置 HPO)。
| 文档 | 用途 |
|---|---|
演进方案.md |
Agentic 契约 / WP / 验收 DoD |
skills/leaves-autotrain/SKILL.md |
决策表与闭环策略 |
skills/leaves-autotrain/cli.md |
flag + metrics schema |
examples/autotrain/ |
零准备 demo |
sniff → train(--cv --runs --tag [--early-stop 默认 save-best] [--from-run])
→ 读 metrics/runs → 决策表调参 → 收敛 → inspect → publish(--emit-repro-script)
边界:本地工件包(非 registry);不内置 HPO;CSV 默认无缺失(--na-policy skip-row 可丢行,不做插补)。
go run ./cmd/leaves --error-format json sniff --data train.csv
go run ./cmd/leaves train --data train.csv --objective reg:squarederror --cv 5 --runs runs.jsonl --tag baseline
go run ./cmd/leaves train --data train.csv --from-run runs.jsonl --tag baseline --out-model m.leaves.json --metrics metrics.json
go run ./cmd/leaves publish --model m.leaves.json --out-dir release/ --quantize --export-xgb --metrics metrics.json --emit-repro-script bothpublish --quantize 会持久化 int8 量化侧车(model.quant.json);manifest.json 含 reproduce 复现命令与文件 sha256。
| 文档 | 说明 |
|---|---|
| godoc | API 参考 |
| docs/api-surface.md | 推荐 / 兼容 / 实验 API 分层与迁移 |
| docs/versioning.md | v2.x 允许改什么 |
| docs/release-checklist.md | v2.1 发版检查表 |
| CHANGELOG.md | 面向用户的变更记录(Unreleased → tag) |
| docs/interop-matrix.md | 格式支持等级 |
| docs/backend-auto.md | BackendAuto 2.0 决策表 |
| docs/extension-points.md | 自定义 objective/metric |
| 演进计划.md | 库 12 个月路线(v5.4) |
| 演进方案.md | Agent 闭环契约(已达成) |
| TODO.md | 可执行 backlog |
| NOTES.md | 仍有效的兼容注记(AutoTransform 等) |
| compatibility.md | 外部 GBRT 库正确性矩阵 |
| AGENTS.md | 项目规约 |
| docs/testdata-matrix.md | 回归矩阵 |
| docs/benchmark-baseline.md | Bench 门禁 |
| examples/ | autotrain / wasm / http / serving-template / train_from_model / extension / multitarget |
- API 分层与迁移:docs/api-surface.md
- 版本策略:docs/versioning.md
- 行为注记(AutoTransform 默认等):NOTES.md
- 外部库版本矩阵:compatibility.md
单线程 / batch=1000,硬件:2017 款 15 寸 MacBook Pro(2.9 GHz Core i7,16 GB 2133 MHz LPDDR3)。C API 走 Python 绑定,批调用下绑定开销可忽略。go test -bench 驱动 leaves 一侧;数据准备流程见 benchmark/ 和 testdata/README.md。
单线程:
| 测试用例 | 特征数 | 树数 | 批量 | C API | leaves |
|---|---|---|---|---|---|
| LightGBM MS LTR | 137 | 500 | 1000 | 49 ms | 51 ms |
| LightGBM Higgs | 28 | 500 | 1000 | 50 ms | 50 ms |
| LightGBM KDD Cup 99* | 41 | 1200 | 1000 | 70 ms | 85 ms |
| XGBoost Higgs | 28 | 500 | 1000 | 44 ms | 50 ms |
四线程:
| 测试用例 | 特征数 | 树数 | 批量 | C API | leaves |
|---|---|---|---|---|---|
| LightGBM MS LTR | 137 | 500 | 1000 | 14 ms | 14 ms |
| LightGBM Higgs | 28 | 500 | 1000 | 14 ms | 14 ms |
| LightGBM KDD Cup 99* | 41 | 1200 | 1000 | 19 ms | 24 ms |
| XGBoost Higgs | 28 | 500 | 1000 | ? | 14 ms |
(?)XGBoost 的多线程预测目前无法通过 Python 绑定启用。(*)KDD Cup 99 同时含连续与类别特征。
Born 后端的微基准(Windows / WebGPU)见 docs/benchmark-baseline.md。
- LightGBM —— 支持的变换函数有限(仅 sigmoid 与 softmax)
- XGBoost —— 支持的变换函数有限(仅 sigmoid 与 softmax);浮点转换的精度差异会导致 C API 与
_leaves_的预测有零星 ULP 级偏差 - scikit-learn —— 不支持变换函数(输出为
GradientBoostingClassifier.decision_function的原始分数);仅支持 pickle 协议0;同样有浮点精度差异
路线图现已转为“当前代码库审计 + 下一年演进计划”的写法。见 演进计划.md 的 v5.4,以及作为历史已完成清单的 TODO.md。
有任何问题或想交流,可发邮件:steper@foxmail.com。
MIT —— Copyright (c) 2018 Steper Lin。
