模型 IO 简介
自 2.1.0 版本以来,XGBoost 的默认模型格式是 UBJSON 格式,此选项适用于模型序列化到文件、模型序列化到缓冲区以及内存快照(pickle 及类似操作)。
JSON 和 UBJSON 具有相同的文档结构,但表示形式不同,我们统称它们为 JSON 格式。本教程旨在分享关于 XGBoost 中使用的 JSON 序列化方法的一些基本见解。除非明确提及,以下部分假设您正在使用这两种输出格式之一,这可以通过在保存/加载模型时提供带有 .json
(或二进制 JSON 的 .ubj
)文件扩展名的文件名来启用:booster.save_model('model.json')
。更多详情请见下文。
在我们开始之前,XGBoost 是一个专注于树模型的梯度提升库,这意味着在 XGBoost 内部,有 2 个不同的部分:
由树组成的模型,以及
用于构建模型的超参数和配置。
如果您来自深度学习社区,那么您应该清楚,由带有固定张量操作的权重组成的神经网络结构与用于训练它们的优化器(如 RMSprop)之间存在差异。
因此,当调用 booster.save_model
(R 中的 xgb.save
)时,XGBoost 会保存树、一些模型参数(例如训练过的树中的输入列数)以及目标函数,它们结合起来代表 XGBoost 中的“模型”概念。至于为什么我们将目标函数作为模型的一部分保存,那是因为目标函数控制着全局偏差(在 XGBoost 中称为 base_score
或截距)和特定于任务的信息的转换。用户可以与他人共享此模型以进行推断、评估或使用不同的超参数集继续训练。
然而,这并非故事的结局。在某些情况下,我们需要保存的不仅仅是模型本身。例如,在分布式训练中,XGBoost 执行检查点操作。或者出于某些原因,您喜欢的分布式计算框架决定将模型从一个工作节点复制到另一个工作节点并在那里继续训练。在这种情况下,序列化输出需要包含足够的信息以继续之前的训练,而无需用户再次提供任何参数。我们将此类场景视为内存快照(或基于内存的序列化方法),并将其与正常的模型 IO 操作区分开来。目前,内存快照用于以下地方:
Python 包:当
Booster
对象使用内置的pickle
模块进行腌制(pickled)时。R 包:当
xgb.Booster
对象使用内置函数saveRDS
或save
进行持久化时。JVM 包:当
Booster
对象使用内置函数saveModel
进行序列化时。
要为模型 IO(仅保存树和目标)启用 JSON 格式支持,请提供一个带有 .json
或 .ubj
作为文件扩展名的文件名,后者是 通用二进制 JSON 的扩展名
bst.save_model('model_file_name.json')
xgb.save(bst, 'model_file_name.json')
val format = "json" // or val format = "ubj"
model.write.option("format", format).save("model_directory_path")
注意
仅从由 XGBoost 生成的 JSON 文件加载模型。尝试加载由外部源生成的 JSON 文件可能会导致未定义行为和崩溃。
当重新加载模型时,XGBoost 会识别文件扩展名 .json
和 .ubj
,并能相应地进行分发。如果未指定扩展名,XGBoost 会尝试猜测正确的扩展名。
关于模型和内存快照的向后兼容性说明
我们保证模型的向后兼容性,但不保证内存快照的向后兼容性。
模型(树和目标)使用稳定的表示,因此早期版本的 XGBoost 生成的模型可以在后续版本的 XGBoost 中访问。如果您想将模型存储或存档以进行长期保存,请使用 save_model
(Python) 和 xgb.save
(R)。
另一方面,内存快照(序列化)捕获了 XGBoost 内部的许多内容,其格式不稳定且经常变化。因此,内存快照仅适用于检查点操作,在这种情况下,您将训练配置的完整快照持久化,以便能够从可能的故障中稳健地恢复并恢复训练过程。加载由早期版本的 XGBoost 生成的内存快照可能会导致错误或未定义行为。如果模型使用 pickle.dump
(Python) 或 saveRDS
(R) 进行持久化,则该模型可能无法在后续版本的 XGBoost 中访问。
自定义目标和指标
XGBoost 接受用户提供的目标和指标函数作为扩展。这些函数不会保存到模型文件中,因为它们是依赖于语言的特性。使用 Python,用户可以通过 pickle 模型将这些函数包含在保存的二进制文件中。一个缺点是,pickle 的输出不是稳定的序列化格式,并且不适用于不同的 Python 版本或 XGBoost 版本,更不用说不同的语言环境了。另一种解决此限制的方法是在加载模型后再次提供这些函数。如果自定义函数很有用,请考虑提交 PR 以在 XGBoost 内部实现它,这样我们的函数就可以与不同的语言绑定一起工作。
从不同版本的 XGBoost 加载 pickled 文件
如前所述,pickled 模型既不可移植也不稳定,但在某些情况下,pickled 模型很有价值。将来恢复它的一种方法是使用特定版本的 Python 和 XGBoost 重新加载它,然后通过调用 save_model 导出模型。
类似的过程可用于恢复旧 RDS 文件中持久化的模型。在 R 中,您可以使用 remotes
包安装旧版本的 XGBoost
library(remotes)
remotes::install_version("xgboost", "0.90.0.1") # Install version 0.90.0.1
安装所需版本后,您可以使用 readRDS
加载 RDS 文件并恢复 xgb.Booster
对象。然后调用 xgb.save
使用稳定表示导出模型。现在您应该能够在最新版本的 XGBoost 中使用该模型。
保存和加载内部参数配置
XGBoost 的 C API
、Python API
和 R API
支持直接将内部配置保存和加载为 JSON 字符串。在 Python 包中
bst = xgboost.train(...)
config = bst.save_config()
print(config)
或在 R 中
config <- xgb.config(bst)
print(config)
将打印出类似于(非实际输出,因为它太长无法演示)
{
"Learner": {
"generic_parameter": {
"device": "cuda:0",
"gpu_page_size": "0",
"n_jobs": "0",
"random_state": "0",
"seed": "0",
"seed_per_iteration": "0"
},
"gradient_booster": {
"gbtree_train_param": {
"num_parallel_tree": "1",
"process_type": "default",
"tree_method": "hist",
"updater": "grow_gpu_hist",
"updater_seq": "grow_gpu_hist"
},
"name": "gbtree",
"updater": {
"grow_gpu_hist": {
"gpu_hist_train_param": {
"debug_synchronize": "0",
},
"train_param": {
"alpha": "0",
"cache_opt": "1",
"colsample_bylevel": "1",
"colsample_bynode": "1",
"colsample_bytree": "1",
"default_direction": "learn",
...
"subsample": "1"
}
}
}
},
"learner_train_param": {
"booster": "gbtree",
"disable_default_eval_metric": "0",
"objective": "reg:squarederror"
},
"metrics": [],
"objective": {
"name": "reg:squarederror",
"reg_loss_param": {
"scale_pos_weight": "1"
}
}
},
"version": [1, 0, 0]
}
您可以将其加载回由相同版本 XGBoost 生成的模型,通过
bst.load_config(config)
这样用户就可以更密切地研究内部表示。请注意,一些 JSON 生成器使用与语言环境相关的浮点序列化方法,XGBoost 不支持这些方法。
保存模型和转储模型之间的区别
XGBoost 在 Booster 类中有一个名为 dump_model
的函数,它允许您以可读格式导出模型,例如 text
、json
或 dot
(graphviz)。它的主要用例是模型解释和可视化,不应重新加载到 XGBoost 中。JSON 版本有一个 schema。更多信息请参见下一节。
JSON 模式
JSON 格式的另一个重要特性是文档化的 模式,基于此可以轻松地重用 XGBoost 的输出模型。这里是输出模型(非序列化,如上所述,它将不稳定)的 JSON 模式。有关解析 XGBoost 树模型的示例,请参见 /demo/json-model
。请注意“dart”增强器中使用的“weight_drop”字段。XGBoost 不直接缩放树叶,而是将权重保存为单独的数组。
{
"$schema": "https://json-schema.fullstack.org.cn/draft-07/schema#",
"definitions": {
"gbtree": {
"type": "object",
"properties": {
"name": {
"const": "gbtree"
},
"model": {
"type": "object",
"properties": {
"gbtree_model_param": {
"$ref": "#/definitions/gbtree_model_param"
},
"trees": {
"type": "array",
"items": {
"type": "object",
"properties": {
"tree_param": {
"$ref": "#/definitions/tree_param"
},
"id": {
"type": "integer"
},
"loss_changes": {
"type": "array",
"items": {
"type": "number"
}
},
"sum_hessian": {
"type": "array",
"items": {
"type": "number"
}
},
"base_weights": {
"type": "array",
"items": {
"type": "number"
}
},
"left_children": {
"type": "array",
"items": {
"type": "integer"
}
},
"right_children": {
"type": "array",
"items": {
"type": "integer"
}
},
"parents": {
"type": "array",
"items": {
"type": "integer"
}
},
"split_indices": {
"type": "array",
"items": {
"type": "integer"
}
},
"split_conditions": {
"type": "array",
"items": {
"type": "number"
}
},
"split_type": {
"type": "array",
"items": {
"type": "integer"
}
},
"default_left": {
"type": "array",
"items": {
"type": "integer"
}
},
"categories": {
"type": "array",
"items": {
"type": "integer"
}
},
"categories_nodes": {
"type": "array",
"items": {
"type": "integer"
}
},
"categories_segments": {
"type": "array",
"items": {
"type": "integer"
}
},
"categories_sizes": {
"type": "array",
"items": {
"type": "integer"
}
}
},
"required": [
"tree_param",
"loss_changes",
"sum_hessian",
"base_weights",
"left_children",
"right_children",
"parents",
"split_indices",
"split_conditions",
"default_left",
"categories",
"categories_nodes",
"categories_segments",
"categories_sizes"
]
}
},
"tree_info": {
"type": "array",
"items": {
"type": "integer"
}
}
},
"required": [
"gbtree_model_param",
"trees",
"tree_info"
]
}
},
"required": [
"name",
"model"
]
},
"gbtree_model_param": {
"type": "object",
"properties": {
"num_trees": {
"type": "string"
},
"num_parallel_tree": {
"type": "string"
}
},
"required": [
"num_trees",
"num_parallel_tree"
]
},
"tree_param": {
"type": "object",
"properties": {
"num_nodes": {
"type": "string"
},
"size_leaf_vector": {
"type": "string"
},
"num_feature": {
"type": "string"
}
},
"required": [
"num_nodes",
"num_feature",
"size_leaf_vector"
]
},
"reg_loss_param": {
"type": "object",
"properties": {
"scale_pos_weight": {
"type": "string"
}
}
},
"pseudo_huber_param": {
"type": "object",
"properties": {
"huber_slope": {
"type": "string"
}
}
},
"aft_loss_param": {
"type": "object",
"properties": {
"aft_loss_distribution": {
"type": "string"
},
"aft_loss_distribution_scale": {
"type": "string"
}
}
},
"softmax_multiclass_param": {
"type": "object",
"properties": {
"num_class": { "type": "string" }
}
},
"lambda_rank_param": {
"type": "object",
"properties": {
"num_pairsample": { "type": "string" },
"fix_list_weight": { "type": "string" }
}
},
"lambdarank_param": {
"type": "object",
"properties": {
"lambdarank_num_pair_per_sample": { "type": "string" },
"lambdarank_pair_method": { "type": "string" },
"lambdarank_unbiased": {"type": "string" },
"lambdarank_bias_norm": {"type": "string" },
"ndcg_exp_gain": {"type": "string"}
}
}
},
"type": "object",
"properties": {
"version": {
"type": "array",
"items": [
{
"type": "number",
"minimum": 1
},
{
"type": "number",
"minimum": 0
},
{
"type": "number",
"minimum": 0
}
],
"minItems": 3,
"maxItems": 3
},
"learner": {
"type": "object",
"properties": {
"feature_names": {
"type": "array",
"items": {
"type": "string"
}
},
"feature_types": {
"type": "array",
"items": {
"type": "string"
}
},
"gradient_booster": {
"oneOf": [
{
"$ref": "#/definitions/gbtree"
},
{
"type": "object",
"properties": {
"name": { "const": "gblinear" },
"model": {
"type": "object",
"properties": {
"weights": {
"type": "array",
"items": {
"type": "number"
}
}
}
}
}
},
{
"type": "object",
"properties": {
"name": { "const": "dart" },
"gbtree": {
"$ref": "#/definitions/gbtree"
},
"weight_drop": {
"type": "array",
"items": {
"type": "number"
}
}
},
"required": [
"name",
"gbtree",
"weight_drop"
]
}
]
},
"objective": {
"oneOf": [
{
"type": "object",
"properties": {
"name": { "const": "reg:squarederror" },
"reg_loss_param": { "$ref": "#/definitions/reg_loss_param"}
},
"required": [
"name",
"reg_loss_param"
]
},
{
"type": "object",
"properties": {
"name": { "const": "reg:pseudohubererror" },
"reg_loss_param": { "$ref": "#/definitions/reg_loss_param"}
},
"required": [
"name",
"reg_loss_param"
]
},
{
"type": "object",
"properties": {
"name": { "const": "reg:squaredlogerror" },
"reg_loss_param": { "$ref": "#/definitions/reg_loss_param"}
},
"required": [
"name",
"reg_loss_param"
]
},
{
"type": "object",
"properties": {
"name": { "const": "reg:linear" },
"reg_loss_param": { "$ref": "#/definitions/reg_loss_param"}
},
"required": [
"name",
"reg_loss_param"
]
},
{
"type": "object",
"properties": {
"name": { "const": "reg:logistic" },
"reg_loss_param": { "$ref": "#/definitions/reg_loss_param"}
},
"required": [
"name",
"reg_loss_param"
]
},
{
"type": "object",
"properties": {
"name": { "const": "binary:logistic" },
"reg_loss_param": { "$ref": "#/definitions/reg_loss_param"}
},
"required": [
"name",
"reg_loss_param"
]
},
{
"type": "object",
"properties": {
"name": { "const": "binary:logitraw" },
"reg_loss_param": { "$ref": "#/definitions/reg_loss_param"}
},
"required": [
"name",
"reg_loss_param"
]
},
{
"type": "object",
"properties": {
"name": { "const": "count:poisson" },
"poisson_regression_param": {
"type": "object",
"properties": {
"max_delta_step": { "type": "string" }
}
}
},
"required": [
"name",
"poisson_regression_param"
]
},
{
"type": "object",
"properties": {
"name": { "const": "reg:tweedie" },
"tweedie_regression_param": {
"type": "object",
"properties": {
"tweedie_variance_power": { "type": "string" }
}
}
},
"required": [
"name",
"tweedie_regression_param"
]
},
{
"properties": {
"name": {
"const": "reg:absoluteerror"
}
},
"type": "object"
},
{
"properties": {
"name": {
"const": "reg:quantileerror"
},
"quantile_loss_param": {
"type": "object",
"properties": {
"quantle_alpha": {"type": "array"}
}
}
},
"type": "object"
},
{
"type": "object",
"properties": {
"name": { "const": "survival:cox" }
},
"required": [ "name" ]
},
{
"type": "object",
"properties": {
"name": { "const": "reg:gamma" }
},
"required": [ "name" ]
},
{
"type": "object",
"properties": {
"name": { "const": "multi:softprob" },
"softmax_multiclass_param": { "$ref": "#/definitions/softmax_multiclass_param"}
},
"required": [
"name",
"softmax_multiclass_param"
]
},
{
"type": "object",
"properties": {
"name": { "const": "multi:softmax" },
"softmax_multiclass_param": { "$ref": "#/definitions/softmax_multiclass_param"}
},
"required": [
"name",
"softmax_multiclass_param"
]
},
{
"type": "object",
"properties": {
"name": { "const": "rank:pairwise" },
"lambda_rank_param": { "$ref": "#/definitions/lambdarank_param"}
},
"required": [
"name",
"lambdarank_param"
]
},
{
"type": "object",
"properties": {
"name": { "const": "rank:ndcg" },
"lambda_rank_param": { "$ref": "#/definitions/lambdarank_param"}
},
"required": [
"name",
"lambdarank_param"
]
},
{
"type": "object",
"properties": {
"name": { "const": "rank:map" },
"lambda_rank_param": { "$ref": "#/definitions/lambda_rank_param"}
},
"required": [
"name",
"lambda_rank_param"
]
},
{
"type": "object",
"properties": {
"name": {"const": "survival:aft"},
"aft_loss_param": { "$ref": "#/definitions/aft_loss_param"}
}
},
{
"type": "object",
"properties": {
"name": {"const": "binary:hinge"}
}
}
]
},
"learner_model_param": {
"type": "object",
"properties": {
"base_score": { "type": "string" },
"num_class": { "type": "string" },
"num_feature": { "type": "string" },
"num_target": { "type": "string" }
}
}
},
"required": [
"gradient_booster",
"objective"
]
}
},
"required": [
"version",
"learner"
]
}
简史
JSON 格式于 1.0 版本引入,旨在用一个可以轻松重用的开放格式取代现已删除的旧二进制内部格式。
后来在 XGBoost 1.6.0 中,引入了对通用二进制 JSON 的额外支持,作为更高效模型 IO 的优化。
UBJSON 已在 2.1 版本中设置为默认格式。
旧的二进制格式已在 3.1 版本中删除。