4. 模型训练与保存:使用 XGBoost 训练一个分类模型、模型序列化保存为 UBJ 格式
本章是部署流水线的核心环节。我们将训练一个 XGBoost 二分类模型,并采用 UBJ (Universal Binary JSON) 格式进行序列化保存。UBJ 格式相比传统的 pickle 或 JSON 具有更快的反序列化速度和更小的体积,非常适合实时评分场景下的模型加载。
4.1 训练数据集与任务定义
我们使用公开的 UCI 信用卡欺诈检测数据集(模拟版本)进行演示。任务目标:根据交易特征预测该笔交易是否为欺诈行为(1 表示欺诈,0 表示正常)。
| 特征组 | 特征名称 | 数据类型 | 说明 |
|---|---|---|---|
| 时间特征 | Time | float | 交易发生时间(秒) |
| 金额特征 | Amount | float | 交易金额 |
| PCA 特征 | V1 ~ V28 | float | 经过 PCA 降维后的匿名特征 |
| 标签 | Class | int (0/1) | 目标变量 |
部署视角提示: 实时评分模型的特征工程必须在训练和推理时保持完全一致。建议将特征名称列表保存为元数据,随模型一起部署。
4.2 模型训练核心参数
针对实时评分场景,我们选择 轻量级树模型 配置,在保证精度的同时控制模型大小和推理延迟。
| 参数名 | 设定值 | 选择理由 |
|---|---|---|
| max_depth | 6 | 控制树深度,防止过拟合,同时保持推理速度 |
| n_estimators | 100 | 基础树数量,平衡精度与模型体积 |
| learning_rate | 0.1 | 步长收缩,提升泛化能力 |
| subsample | 0.8 | 行采样,增强鲁棒性 |
| colsample_bytree | 0.8 | 列采样,减少过拟合 |
| eval_metric | auc | 适合不平衡分类场景 |
| objective | binary:logistic | 二分类输出概率 |
4.3 训练代码实现
以下代码展示了完整的训练流程,包括数据加载、参数设置、模型训练和早停机制。
import xgboost as xgb
import pandas as pd
from sklearn.model_selection import train_test_split
# 1. 加载数据
df = pd.read_csv('creditcard.csv')
X = df.drop('Class', axis=1)
y = df['Class']
# 2. 划分训练集和验证集
X_train, X_val, y_train, y_val = train_test_split(
X, y, test_size=0.2, random_state=42, stratify=y
)
# 3. 转换为 DMatrix 格式(XGBoost 原生数据结构)
dtrain = xgb.DMatrix(X_train, label=y_train)
dval = xgb.DMatrix(X_val, label=y_val)
# 4. 参数配置
params = {
'max_depth': 6,
'learning_rate': 0.1,
'n_estimators': 100,
'subsample': 0.8,
'colsample_bytree': 0.8,
'eval_metric': 'auc',
'objective': 'binary:logistic',
'seed': 42
}
# 5. 训练(带早停)
model = xgb.train(
params,
dtrain,
num_boost_round=100,
evals=[(dtrain, 'train'), (dval, 'val')],
early_stopping_rounds=10,
verbose_eval=10
)
print(f"最佳迭代轮数: {model.best_iteration}")
print(f"最佳 AUC: {model.best_score}")
关键点: 使用
early_stopping_rounds 可以自动确定最优树数量,避免手动调参。训练完成后,model.best_iteration 记录了最佳轮次。
4.4 模型序列化:UBJ 格式
UBJ (Universal Binary JSON) 是 XGBoost 原生支持的二进制序列化格式。相比 JSON 文本格式,UBJ 具有以下优势:
- 反序列化速度提升 3-5 倍:二进制解析无需字符到数值的转换。
- 体积减少 40%-60%:无冗余空格和引号。
- 跨语言兼容:C++、Python、Java、R 等语言均可直接加载。
保存与加载的代码实现:
import xgboost as xgb
# 保存为 UBJ 格式
model.save_model('fraud_detection_model.ubj')
# 加载 UBJ 模型(用于验证)
loaded_model = xgb.Booster()
loaded_model.load_model('fraud_detection_model.ubj')
# 验证加载后的模型与原始模型预测一致
import numpy as np
pred_original = model.predict(dval)
pred_loaded = loaded_model.predict(dval)
assert np.allclose(pred_original, pred_loaded), "模型加载前后预测不一致!"
print("UBJ 模型保存与加载验证通过。")
注意事项:
- UBJ 格式的文件扩展名通常使用
.ubj,但 XGBoost 不强制要求。 - 保存时不要使用
model.save_model('model.json')的 JSON 格式,因为 JSON 在实时评分场景下反序列化开销较大。 - 生产环境中建议在保存后立即执行加载验证,确保模型文件未损坏。
4.5 模型训练与保存流程总览
下图展示了从原始数据到 UBJ 模型文件的完整工作流:
4.6 模型文件结构说明
UBJ 格式的模型文件内部采用 自描述二进制结构,包含以下核心组件:
| 组件 | 说明 | 大小占比 |
|---|---|---|
| 模型头 | 版本号、目标函数、基学习器类型等元信息 | ~1% |
| 树结构 | 每棵树的节点分裂特征、阈值、叶子权重 | ~85% |
| 参数配置 | 训练时使用的超参数(用于推理一致性) | ~5% |
| 特征映射 | 特征名称到索引的映射表(可选) | ~9% |
部署优化建议: 如果模型文件过大(超过 100MB),可以考虑:
1. 降低
2. 使用
3. 在推理服务启动时预加载模型到内存,避免每次请求都读取磁盘。
1. 降低
max_depth 或 n_estimators。2. 使用
model.save_model('model.ubj', compression='gzip') 启用压缩。3. 在推理服务启动时预加载模型到内存,避免每次请求都读取磁盘。
4.7 本章小结
- 使用 XGBoost 训练二分类模型时,早停机制 可以有效防止过拟合并自动确定最优树数量。
- UBJ 格式 是实时评分场景下的首选序列化方案,兼具速度和体积优势。
- 保存模型后必须执行 加载验证,确保模型文件在生产环境中可用。
- 模型文件的大小直接影响服务启动时间和内存占用,需根据部署环境合理调整训练参数。