Integration tests for interaction constraints with column-wise data split (#9611)

This commit is contained in:
Rong Ou 2023-09-26 17:27:43 -07:00 committed by GitHub
parent c75a3bc0a9
commit d6d14d0fb9
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -720,48 +720,39 @@ INSTANTIATE_TEST_SUITE_P(ColumnSplitObjective, TestColumnSplit,
});
namespace {
void VerifyColumnSplitColumnSampler(std::string const& tree_method, bool use_gpu,
Json const& expected_model) {
Json GetModelWithArgs(std::shared_ptr<DMatrix> dmat, std::string const& tree_method,
std::string const& device, Args const& args) {
std::unique_ptr<Learner> learner{Learner::Create({dmat})};
learner->SetParam("tree_method", tree_method);
learner->SetParam("device", device);
learner->SetParam("objective", "reg:logistic");
learner->SetParams(args);
learner->UpdateOneIter(0, dmat);
Json model{Object{}};
{
auto const world_size = collective::GetWorldSize();
auto const rank = collective::GetRank();
auto const objective = "reg:logistic";
auto p_fmat = MakeFmatForObjTest(objective);
std::shared_ptr<DMatrix> sliced{p_fmat->SliceCol(world_size, rank)};
std::unique_ptr<Learner> learner{Learner::Create({sliced})};
learner->SetParam("tree_method", tree_method);
if (use_gpu) {
auto gpu_id = common::AllVisibleGPUs() == 1 ? 0 : rank;
learner->SetParam("device", "cuda:" + std::to_string(gpu_id));
}
learner->SetParam("objective", objective);
learner->SetParam("colsample_bytree", "0.5");
learner->SetParam("colsample_bylevel", "0.6");
learner->SetParam("colsample_bynode", "0.7");
learner->UpdateOneIter(0, sliced);
learner->SaveModel(&model);
learner->SaveModel(&model);
return model;
}
void VerifyColumnSplitWithArgs(std::string const& tree_method, bool use_gpu, Args const& args,
Json const& expected_model) {
auto const world_size = collective::GetWorldSize();
auto const rank = collective::GetRank();
auto p_fmat = MakeFmatForObjTest("");
std::shared_ptr<DMatrix> sliced{p_fmat->SliceCol(world_size, rank)};
std::string device = "cpu";
if (use_gpu) {
auto gpu_id = common::AllVisibleGPUs() == 1 ? 0 : rank;
device = "cuda:" + std::to_string(gpu_id);
}
auto model = GetModelWithArgs(sliced, tree_method, device, args);
ASSERT_EQ(model, expected_model);
}
void TestColumnSplitColumnSampler(std::string const& tree_method, bool use_gpu) {
Json model{Object{}};
{
auto objective = "reg:logistic";
auto p_fmat = MakeFmatForObjTest(objective);
std::unique_ptr<Learner> learner{Learner::Create({p_fmat})};
learner->SetParam("tree_method", tree_method);
if (use_gpu) {
learner->SetParam("device", "cuda:0");
}
learner->SetParam("objective", objective);
learner->SetParam("colsample_bytree", "0.5");
learner->SetParam("colsample_bylevel", "0.6");
learner->SetParam("colsample_bynode", "0.7");
learner->UpdateOneIter(0, p_fmat);
learner->SaveModel(&model);
}
void TestColumnSplitWithArgs(std::string const& tree_method, bool use_gpu, Args const& args) {
auto p_fmat = MakeFmatForObjTest("");
std::string device = use_gpu ? "cuda:0" : "cpu";
auto model = GetModelWithArgs(p_fmat, tree_method, device, args);
auto world_size{3};
if (use_gpu) {
world_size = common::AllVisibleGPUs();
@ -770,9 +761,19 @@ void TestColumnSplitColumnSampler(std::string const& tree_method, bool use_gpu)
world_size = 3;
}
}
RunWithInMemoryCommunicator(world_size, VerifyColumnSplitColumnSampler, tree_method, use_gpu,
RunWithInMemoryCommunicator(world_size, VerifyColumnSplitWithArgs, tree_method, use_gpu, args,
model);
}
void TestColumnSplitColumnSampler(std::string const& tree_method, bool use_gpu) {
Args args{{"colsample_bytree", "0.5"}, {"colsample_bylevel", "0.6"}, {"colsample_bynode", "0.7"}};
TestColumnSplitWithArgs(tree_method, use_gpu, args);
}
void TestColumnSplitInteractionConstraints(std::string const& tree_method, bool use_gpu) {
Args args{{"interaction_constraints", "[[0, 5, 7], [2, 8, 9], [1, 3, 6]]"}};
TestColumnSplitWithArgs(tree_method, use_gpu, args);
}
} // anonymous namespace
TEST(ColumnSplitColumnSampler, Approx) { TestColumnSplitColumnSampler("approx", false); }
@ -780,8 +781,26 @@ TEST(ColumnSplitColumnSampler, Approx) { TestColumnSplitColumnSampler("approx",
TEST(ColumnSplitColumnSampler, Hist) { TestColumnSplitColumnSampler("hist", false); }
#if defined(XGBOOST_USE_CUDA)
TEST(ColumnSplitColumnSampler, GPUApprox) { TestColumnSplitColumnSampler("approx", true); }
TEST(MGPUColumnSplitColumnSampler, GPUApprox) { TestColumnSplitColumnSampler("approx", true); }
TEST(ColumnSplitColumnSampler, GPUHist) { TestColumnSplitColumnSampler("hist", true); }
TEST(MGPUColumnSplitColumnSampler, GPUHist) { TestColumnSplitColumnSampler("hist", true); }
#endif // defined(XGBOOST_USE_CUDA)
TEST(ColumnSplitInteractionConstraints, Approx) {
TestColumnSplitInteractionConstraints("approx", false);
}
TEST(ColumnSplitInteractionConstraints, Hist) {
TestColumnSplitInteractionConstraints("hist", false);
}
#if defined(XGBOOST_USE_CUDA)
TEST(MGPUColumnSplitInteractionConstraints, GPUApprox) {
TestColumnSplitInteractionConstraints("approx", true);
}
TEST(MGPUColumnSplitInteractionConstraints, GPUHist) {
TestColumnSplitInteractionConstraints("hist", true);
}
#endif // defined(XGBOOST_USE_CUDA)
} // namespace xgboost