Skip to content

Commit

Permalink
Add tests for tree grow policy. (#7302)
Browse files Browse the repository at this point in the history
  • Loading branch information
trivialfis authored Oct 12, 2021
1 parent 5b17bb0 commit 130df8c
Showing 1 changed file with 80 additions and 0 deletions.
80 changes: 80 additions & 0 deletions tests/cpp/tree/test_tree_policy.cc
Original file line number Diff line number Diff line change
@@ -0,0 +1,80 @@
/*!
* Copyright 2021 XGBoost contributors
*/
#include <gtest/gtest.h>
#include <xgboost/base.h>
#include <xgboost/tree_model.h>
#include "../helpers.h"

namespace xgboost {
class TestGrowPolicy : public ::testing::Test {
protected:
std::shared_ptr<DMatrix> Xy_;
size_t n_samples_ = 4096, n_features_ = 13;
float sparsity_ = 0.5;

protected:
void SetUp() override {
Xy_ =
RandomDataGenerator{n_samples_, n_features_, sparsity_}.GenerateDMatrix(
true);
}

void TestTreeGrowPolicy(std::string tree_method, std::string policy) {
{
std::unique_ptr<Learner> learner{Learner::Create({this->Xy_})};
learner->SetParam("tree_method", tree_method);
learner->SetParam("max_leaves", "16");
learner->SetParam("grow_policy", policy);
learner->Configure();

learner->UpdateOneIter(0, Xy_);
Json model{Object{}};
learner->SaveModel(&model);

auto j_tree = model["learner"]["gradient_booster"]["model"]["trees"][0];
RegTree tree;
tree.LoadModel(j_tree);
ASSERT_EQ(tree.GetNumLeaves(), 16);
}
{
std::unique_ptr<Learner> learner{Learner::Create({this->Xy_})};
learner->SetParam("tree_method", tree_method);
learner->SetParam("max_depth", "3");
learner->SetParam("grow_policy", policy);
learner->Configure();

learner->UpdateOneIter(0, Xy_);
Json model{Object{}};
learner->SaveModel(&model);

auto j_tree = model["learner"]["gradient_booster"]["model"]["trees"][0];
RegTree tree;
tree.LoadModel(j_tree);
bst_node_t depth = 0;
tree.WalkTree([&](bst_node_t nidx) {
depth = std::max(tree.GetDepth(nidx), depth);
return true;
});
ASSERT_EQ(depth, 3);
}
}
};

TEST_F(TestGrowPolicy, DISABLED_Approx) {
this->TestTreeGrowPolicy("approx", "depthwise");
this->TestTreeGrowPolicy("approx", "lossguide");
}

TEST_F(TestGrowPolicy, Hist) {
this->TestTreeGrowPolicy("hist", "depthwise");
this->TestTreeGrowPolicy("hist", "lossguide");
}

#if defined(XGBOOST_USE_CUDA)
TEST_F(TestGrowPolicy, GpuHist) {
this->TestTreeGrowPolicy("gpu_hist", "depthwise");
this->TestTreeGrowPolicy("gpu_hist", "lossguide");
}
#endif // defined(XGBOOST_USE_CUDA)
} // namespace xgboost

0 comments on commit 130df8c

Please sign in to comment.