Hendrik/xgboost
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Enforce tree order in JSON. (#5974)

Browse Source 
* Make JSON model IO more future proof by using tree id in model loading.
This commit is contained in:
Jiaming Yuan
2020-08-05 16:44:52 +08:00
committed by GitHub
parent dde9c5aaff
commit 9c6e791e64
2 changed files with 23 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
src/gbm/gbtree_model.cc
View File

@@ -1,6 +1,8 @@
/*!
* Copyright 2019 by Contributors
* Copyright 2019-2020 by Contributors
*/
#include <utility>
#include "xgboost/json.h"
#include "xgboost/logging.h"
#include "gbtree_model.h"
@@ -41,15 +43,14 @@ void GBTreeModel::SaveModel(Json* p_out) const {
auto& out = *p_out;
CHECK_EQ(param.num_trees, static_cast<int>(trees.size()));
out["gbtree_model_param"] = ToJson(param);
std::vector<Json> trees_json;
size_t t = 0;
for (auto const& tree : trees) {
std::vector<Json> trees_json(trees.size());
for (size_t t = 0; t < trees.size(); ++t) {
auto const& tree = trees[t];
Json tree_json{Object()};
tree->SaveModel(&tree_json);
// The field is not used in XGBoost, but might be useful for external project.
tree_json["id"] = Integer(t);
trees_json.emplace_back(tree_json);
t++;
tree_json["id"] = Integer(static_cast<Integer::Int>(t));
trees_json[t] = std::move(tree_json);
}
std::vector<Json> tree_info_json(tree_info.size());
@@ -70,9 +71,10 @@ void GBTreeModel::LoadModel(Json const& in) {
auto const& trees_json = get<Array const>(in["trees"]);
trees.resize(trees_json.size());
for (size_t t = 0; t < trees.size(); ++t) {
trees[t].reset( new RegTree() );
trees[t]->LoadModel(trees_json[t]);
for (size_t t = 0; t < trees_json.size(); ++t) { // NOLINT
auto tree_id = get<Integer>(trees_json[t]["id"]);
trees.at(tree_id).reset(new RegTree());
trees.at(tree_id)->LoadModel(trees_json[t]);
}
tree_info.resize(param.num_trees);