ff26cd3212
Added some more tests for the learner and fit_stump, for both column-wise distributed learning and vertical federated learning. Also moved the `IsRowSplit` and `IsColumnSplit` methods from the `DMatrix` to the `MetaInfo` since in some places we only have access to the `MetaInfo`. Added a new convenience method `IsVerticalFederatedLearning`. Some refactoring of the testing fixtures.
79 lines
2.9 KiB
C++
79 lines
2.9 KiB
C++
/*!
|
|
* Copyright 2023 XGBoost contributors
|
|
*/
|
|
#include <dmlc/parameter.h>
|
|
#include <gtest/gtest.h>
|
|
#include <xgboost/data.h>
|
|
#include <xgboost/objective.h>
|
|
|
|
#include "../../../plugin/federated/federated_server.h"
|
|
#include "../../../src/collective/communicator-inl.h"
|
|
#include "../helpers.h"
|
|
#include "helpers.h"
|
|
|
|
namespace xgboost {
|
|
|
|
class FederatedLearnerTest : public BaseFederatedTest {
|
|
protected:
|
|
static auto constexpr kRows{16};
|
|
static auto constexpr kCols{16};
|
|
};
|
|
|
|
void VerifyBaseScore(size_t rows, size_t cols, float expected_base_score) {
|
|
auto const world_size = collective::GetWorldSize();
|
|
auto const rank = collective::GetRank();
|
|
std::shared_ptr<DMatrix> Xy_{RandomDataGenerator{rows, cols, 0}.GenerateDMatrix(rank == 0)};
|
|
std::shared_ptr<DMatrix> sliced{Xy_->SliceCol(world_size, rank)};
|
|
std::unique_ptr<Learner> learner{Learner::Create({sliced})};
|
|
learner->SetParam("tree_method", "approx");
|
|
learner->SetParam("objective", "binary:logistic");
|
|
learner->UpdateOneIter(0, sliced);
|
|
Json config{Object{}};
|
|
learner->SaveConfig(&config);
|
|
auto base_score = GetBaseScore(config);
|
|
ASSERT_EQ(base_score, expected_base_score);
|
|
}
|
|
|
|
void VerifyModel(size_t rows, size_t cols, Json const& expected_model) {
|
|
auto const world_size = collective::GetWorldSize();
|
|
auto const rank = collective::GetRank();
|
|
std::shared_ptr<DMatrix> Xy_{RandomDataGenerator{rows, cols, 0}.GenerateDMatrix(rank == 0)};
|
|
std::shared_ptr<DMatrix> sliced{Xy_->SliceCol(world_size, rank)};
|
|
std::unique_ptr<Learner> learner{Learner::Create({sliced})};
|
|
learner->SetParam("tree_method", "approx");
|
|
learner->SetParam("objective", "binary:logistic");
|
|
learner->UpdateOneIter(0, sliced);
|
|
Json model{Object{}};
|
|
learner->SaveModel(&model);
|
|
ASSERT_EQ(model, expected_model);
|
|
}
|
|
|
|
TEST_F(FederatedLearnerTest, BaseScore) {
|
|
std::shared_ptr<DMatrix> Xy_{RandomDataGenerator{kRows, kCols, 0}.GenerateDMatrix(true)};
|
|
std::unique_ptr<Learner> learner{Learner::Create({Xy_})};
|
|
learner->SetParam("tree_method", "approx");
|
|
learner->SetParam("objective", "binary:logistic");
|
|
learner->UpdateOneIter(0, Xy_);
|
|
Json config{Object{}};
|
|
learner->SaveConfig(&config);
|
|
auto base_score = GetBaseScore(config);
|
|
ASSERT_NE(base_score, ObjFunction::DefaultBaseScore());
|
|
|
|
RunWithFederatedCommunicator(kWorldSize, server_address_, &VerifyBaseScore, kRows, kCols,
|
|
base_score);
|
|
}
|
|
|
|
TEST_F(FederatedLearnerTest, Model) {
|
|
std::shared_ptr<DMatrix> Xy_{RandomDataGenerator{kRows, kCols, 0}.GenerateDMatrix(true)};
|
|
std::unique_ptr<Learner> learner{Learner::Create({Xy_})};
|
|
learner->SetParam("tree_method", "approx");
|
|
learner->SetParam("objective", "binary:logistic");
|
|
learner->UpdateOneIter(0, Xy_);
|
|
Json model{Object{}};
|
|
learner->SaveModel(&model);
|
|
|
|
RunWithFederatedCommunicator(kWorldSize, server_address_, &VerifyModel, kRows, kCols,
|
|
std::cref(model));
|
|
}
|
|
} // namespace xgboost
|