104 lines
3.3 KiB
C++
104 lines
3.3 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 {
|
|
|
|
void VerifyObjectives(size_t rows, size_t cols, std::vector<float> const &expected_base_scores,
|
|
std::vector<Json> const &expected_models) {
|
|
auto const world_size = collective::GetWorldSize();
|
|
auto const rank = collective::GetRank();
|
|
std::shared_ptr<DMatrix> dmat{RandomDataGenerator{rows, cols, 0}.GenerateDMatrix(rank == 0)};
|
|
|
|
if (rank == 0) {
|
|
auto &h_upper = dmat->Info().labels_upper_bound_.HostVector();
|
|
auto &h_lower = dmat->Info().labels_lower_bound_.HostVector();
|
|
h_lower.resize(rows);
|
|
h_upper.resize(rows);
|
|
for (size_t i = 0; i < rows; ++i) {
|
|
h_lower[i] = 1;
|
|
h_upper[i] = 10;
|
|
}
|
|
}
|
|
std::shared_ptr<DMatrix> sliced{dmat->SliceCol(world_size, rank)};
|
|
|
|
auto i = 0;
|
|
for (auto const *entry : ::dmlc::Registry<::xgboost::ObjFunctionReg>::List()) {
|
|
std::unique_ptr<Learner> learner{Learner::Create({sliced})};
|
|
learner->SetParam("tree_method", "approx");
|
|
learner->SetParam("objective", entry->name);
|
|
if (entry->name.find("quantile") != std::string::npos) {
|
|
learner->SetParam("quantile_alpha", "0.5");
|
|
}
|
|
if (entry->name.find("multi") != std::string::npos) {
|
|
learner->SetParam("num_class", "3");
|
|
}
|
|
learner->UpdateOneIter(0, sliced);
|
|
|
|
Json config{Object{}};
|
|
learner->SaveConfig(&config);
|
|
auto base_score = GetBaseScore(config);
|
|
ASSERT_EQ(base_score, expected_base_scores[i]);
|
|
|
|
Json model{Object{}};
|
|
learner->SaveModel(&model);
|
|
ASSERT_EQ(model, expected_models[i]);
|
|
|
|
i++;
|
|
}
|
|
}
|
|
|
|
class FederatedLearnerTest : public BaseFederatedTest {
|
|
protected:
|
|
static auto constexpr kRows{16};
|
|
static auto constexpr kCols{16};
|
|
};
|
|
|
|
TEST_F(FederatedLearnerTest, Objectives) {
|
|
std::shared_ptr<DMatrix> dmat{RandomDataGenerator{kRows, kCols, 0}.GenerateDMatrix(true)};
|
|
|
|
auto &h_upper = dmat->Info().labels_upper_bound_.HostVector();
|
|
auto &h_lower = dmat->Info().labels_lower_bound_.HostVector();
|
|
h_lower.resize(kRows);
|
|
h_upper.resize(kRows);
|
|
for (size_t i = 0; i < kRows; ++i) {
|
|
h_lower[i] = 1;
|
|
h_upper[i] = 10;
|
|
}
|
|
|
|
std::vector<float> base_scores;
|
|
std::vector<Json> models;
|
|
for (auto const *entry : ::dmlc::Registry<::xgboost::ObjFunctionReg>::List()) {
|
|
std::unique_ptr<Learner> learner{Learner::Create({dmat})};
|
|
learner->SetParam("tree_method", "approx");
|
|
learner->SetParam("objective", entry->name);
|
|
if (entry->name.find("quantile") != std::string::npos) {
|
|
learner->SetParam("quantile_alpha", "0.5");
|
|
}
|
|
if (entry->name.find("multi") != std::string::npos) {
|
|
learner->SetParam("num_class", "3");
|
|
}
|
|
learner->UpdateOneIter(0, dmat);
|
|
Json config{Object{}};
|
|
learner->SaveConfig(&config);
|
|
base_scores.emplace_back(GetBaseScore(config));
|
|
|
|
Json model{Object{}};
|
|
learner->SaveModel(&model);
|
|
models.emplace_back(model);
|
|
}
|
|
|
|
RunWithFederatedCommunicator(kWorldSize, server_address_, &VerifyObjectives, kRows, kCols,
|
|
base_scores, models);
|
|
}
|
|
} // namespace xgboost
|