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

Fix node reuse. (#4404)

Browse Source 
* Reinitialize `_sindex` when reallocating a deleted node.
This commit is contained in:
Jiaming Yuan 2019-04-27 13:03:23 +08:00 committed by GitHub
parent 37dc82c3ff
commit 77c03538b0
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 24 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
include/xgboost/tree_model.h
View File

@ -179,6 +179,10 @@ class RegTree {
XGBOOST_DEVICE void MarkDelete() {
this->sindex_ = std::numeric_limits<unsigned>::max();
}
/*! \brief Reuse this deleted node. */
XGBOOST_DEVICE void Reuse() {
this->sindex_ = 0;
}
// set parent
XGBOOST_DEVICE void SetParent(int pidx, bool is_left_child = true) {
if (is_left_child) pidx |= (1U << 31);
@ -503,10 +507,11 @@ class RegTree {
// !!!!!! NOTE: may cause BUG here, nodes.resize
int AllocNode() {
if (param.num_deleted != 0) {
int nd = deleted_nodes_.back();
int nid = deleted_nodes_.back();
deleted_nodes_.pop_back();
nodes_[nid].Reuse();
--param.num_deleted;
return nd;
return nid;
}
int nd = param.num_nodes++;
CHECK_LT(param.num_nodes, std::numeric_limits<int>::max())

17
tests/cpp/tree/test_tree_model.cc
View File

@ -84,4 +84,21 @@ TEST(Tree, Load) {
EXPECT_EQ(tree[1].LeafValue(), 0.1f);
EXPECT_TRUE(tree[1].IsLeaf());
}
TEST(Tree, AllocateNode) {
RegTree tree;
tree.ExpandNode(
0, 0, 0.0f, false, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f);
tree.CollapseToLeaf(0, 0);
ASSERT_EQ(tree.NumExtraNodes(), 0);
tree.ExpandNode(
0, 0, 0.0f, false, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f);
ASSERT_EQ(tree.NumExtraNodes(), 2);
auto& nodes = tree.GetNodes();
ASSERT_FALSE(nodes.at(1).IsDeleted());
ASSERT_TRUE(nodes.at(1).IsLeaf());
ASSERT_TRUE(nodes.at(2).IsLeaf());
}
} // namespace xgboost