74 lines
2.7 KiB
C++
74 lines
2.7 KiB
C++
/**
|
|
* Copyright 2024, XGBoost Contributors
|
|
*/
|
|
#include <xgboost/base.h>
|
|
#include <xgboost/context.h>
|
|
|
|
#include "../../../src/tree/param.h" // for TrainParam
|
|
#include "../helpers.h" // for RandomDataGenerator
|
|
|
|
namespace xgboost::data {
|
|
template <typename Page, typename Equal>
|
|
void TestExtMemQdmBasic(Context const* ctx, bool on_host, float sparsity, Equal&& check_equal) {
|
|
bst_idx_t n_samples = 256, n_features = 16, n_batches = 4;
|
|
bst_bin_t max_bin = 64;
|
|
bst_target_t n_targets = 3;
|
|
BatchParam p{max_bin, tree::TrainParam::DftSparseThreshold()};
|
|
|
|
auto p_fmat = RandomDataGenerator{n_samples, n_features, sparsity}
|
|
.Bins(max_bin)
|
|
.Batches(n_batches)
|
|
.Targets(n_targets)
|
|
.Device(ctx->Device())
|
|
.OnHost(on_host)
|
|
.GenerateExtMemQuantileDMatrix("temp", true);
|
|
ASSERT_FALSE(p_fmat->SingleColBlock());
|
|
|
|
// Loop over the batches and count the number of pages
|
|
bst_idx_t batch_cnt = 0, base_cnt = 0, row_cnt = 0;
|
|
for (auto const& page : p_fmat->GetBatches<Page>(ctx, p)) {
|
|
ASSERT_EQ(page.BaseRowId(), base_cnt);
|
|
++batch_cnt;
|
|
base_cnt += n_samples / n_batches;
|
|
row_cnt += page.Size();
|
|
ASSERT_EQ((sparsity == 0.0f), page.IsDense());
|
|
}
|
|
ASSERT_EQ(n_batches, batch_cnt);
|
|
ASSERT_EQ(p_fmat->Info().num_row_, n_samples);
|
|
EXPECT_EQ(p_fmat->Info().num_row_, row_cnt);
|
|
ASSERT_EQ(p_fmat->Info().num_col_, n_features);
|
|
if (sparsity == 0.0f) {
|
|
ASSERT_EQ(p_fmat->Info().num_nonzero_, n_samples * n_features);
|
|
} else {
|
|
ASSERT_LT(p_fmat->Info().num_nonzero_, n_samples * n_features);
|
|
ASSERT_GT(p_fmat->Info().num_nonzero_, 0);
|
|
}
|
|
ASSERT_EQ(p_fmat->Info().labels.Shape(0), n_samples);
|
|
ASSERT_EQ(p_fmat->Info().labels.Shape(1), n_targets);
|
|
|
|
// Compare against the sparse page DMatrix
|
|
auto p_sparse = RandomDataGenerator{n_samples, n_features, sparsity}
|
|
.Bins(max_bin)
|
|
.Batches(n_batches)
|
|
.Targets(n_targets)
|
|
.Device(ctx->Device())
|
|
.OnHost(on_host)
|
|
.GenerateSparsePageDMatrix("temp", true);
|
|
auto it = p_fmat->GetBatches<Page>(ctx, p).begin();
|
|
for (auto const& page : p_sparse->GetBatches<Page>(ctx, p)) {
|
|
auto orig = it.Page();
|
|
check_equal(ctx, *orig, page);
|
|
++it;
|
|
}
|
|
|
|
// Check meta info
|
|
auto h_y_sparse = p_sparse->Info().labels.HostView();
|
|
auto h_y = p_fmat->Info().labels.HostView();
|
|
for (std::size_t i = 0, m = h_y_sparse.Shape(0); i < m; ++i) {
|
|
for (std::size_t j = 0, n = h_y_sparse.Shape(1); j < n; ++j) {
|
|
ASSERT_EQ(h_y(i, j), h_y_sparse(i, j));
|
|
}
|
|
}
|
|
}
|
|
} // namespace xgboost::data
|