模型输入输出简介

自2.1.0版本起,XGBoost的默认模型格式为UBJSON格式,该选项已启用以支持将模型序列化到文件、序列化到缓冲区以及内存快照(如pickle等)。

在 XGBoost 1.0.0 版本中,我们引入了对使用 JSON 来保存/加载 XGBoost 模型及相关训练超参数的支持,旨在用一种开放的格式替代旧的二进制内部格式,以便于轻松重用。随后在 XGBoost 1.6.0 版本中,增加了对 Universal Binary JSON 的额外支持,作为优化模型输入输出的更高效方式,该优化在 2.1 版本中被设为默认。

JSON 和 UBJSON 具有相同的文档结构,但表示方式不同,我们将它们统称为 JSON 格式。本教程旨在分享一些关于 XGBoost 中使用的 JSON 序列化方法的基本见解。如果没有特别说明,以下章节假设您使用的是这两种输出格式之一,可以通过在保存/加载模型时提供带有 .json``(或 ``.ubj 用于二进制 JSON)的文件扩展名来启用:booster.save_model('model.json')。更多详情如下。

在我们开始之前,XGBoost 是一个专注于树模型的梯度提升库,这意味着在 XGBoost 内部,有两个不同的部分:

  1. 由树木组成的模型和

  2. 用于构建模型的超参数和配置。

如果你来自深度学习社区,那么你应该清楚,由固定张量操作组成的神经网络结构与用于训练它们的优化器(如RMSprop)之间存在差异。

因此,当调用 booster.save_model``(在 R 中为 ``xgb.save)时,XGBoost 会保存树、一些模型参数(如训练树中的输入列数)以及目标函数,这些共同代表了 XGBoost 中的“模型”概念。至于为什么我们要将目标函数作为模型的一部分保存,那是因为目标函数控制全局偏差(在 XGBoost 中称为 base_score)和任务特定信息的转换。用户可以将此模型与他人共享,用于预测、评估或使用不同的超参数集继续训练等。

然而,这并不是故事的结局。有些情况下,我们需要保存的不仅仅是模型本身。例如,在分布式训练中,XGBoost 执行检查点操作。或者出于某些原因,您喜欢的分布式计算框架决定将模型从一个工作节点复制到另一个工作节点并在那里继续训练。在这种情况下,序列化输出需要包含足够的信息,以便在没有用户再次提供任何参数的情况下继续之前的训练。我们将这种场景视为 **内存快照**(或基于内存的序列化方法),并将其与正常的模型IO操作区分开来。目前,内存快照用于以下地方:

  • Python 包:当 Booster 对象使用内置的 pickle 模块进行序列化时。

  • R 包:当 xgb.Booster 对象通过内置函数 saveRDSsave 持久化时。

  • JVM 包:当 Booster 对象使用内置函数 saveModel 进行序列化时。

其他语言绑定仍在进行中。

备注

旧的二进制格式不区分模型和原始内存序列化格式之间的差异,它是所有内容的混合体,这也是我们希望用更健壮的序列化方法替换它的部分原因。JVM 包有自己的基于内存的序列化方法。

要为模型IO(仅保存树和目标)启用JSON格式支持,请提供一个以``.json``或``.ubj``为文件扩展名的文件名,后者是`Universal Binary JSON <https://ubjson.org/>`__的扩展名。

Python
bst.save_model('model_file_name.json')
xgb.save(bst, 'model_file_name.json')
Scala
val format = "json"  // or val format = "ubj"
model.write.option("format", format).save("model_directory_path")

备注

仅从由 XGBoost 生成的 JSON 文件中加载模型。尝试加载由外部源生成的 JSON 文件可能导致未定义行为和崩溃。

对于内存快照,从 xgboost 1.6 开始,UBJSON 是默认的格式。当加载模型时,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 加载序列化文件

如前所述,序列化的模型既不具有可移植性也不稳定,但在某些情况下,序列化的模型是有价值的。未来恢复它的一个方法是使用特定版本的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 APIPython APIR 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 的函数,它允许你以 textjson``dot``(graphviz)等可读格式导出模型。其主要用途是用于模型解释或可视化,而不是为了重新加载回 XGBoost。JSON 版本有一个 schema。更多信息请参见下一节。

JSON Schema

JSON 格式的另一个重要特性是文档化的 schema ,基于此可以轻松重用 XGBoost 的输出模型。以下是输出模型的 JSON schema(非序列化,如上所述,序列化将不稳定)。有关解析 XGBoost 树模型的示例,请参见 /demo/json-model。请注意在“dart”增强器中使用的“weight_drop”字段。XGBoost 不直接缩放树叶子,而是将权重保存为单独的数组。

{
  "$schema": "http://json-schema.org/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"
  ]
}