Support categorical data in GPU weighted sketching. (#6508)

This commit is contained in:
Jiaming Yuan
2020-12-16 14:23:28 +08:00
committed by GitHub
parent 5c8ccf4455
commit 886486a519
2 changed files with 21 additions and 6 deletions

View File

@@ -122,10 +122,21 @@ TEST(HistUtil, DeviceSketchCategoricalAsNumeric) {
}
}
void TestCategoricalSketch(size_t n, size_t num_categories, int32_t num_bins) {
void TestCategoricalSketch(size_t n, size_t num_categories, int32_t num_bins, bool weighted) {
auto x = GenerateRandomCategoricalSingleColumn(n, num_categories);
auto dmat = GetDMatrixFromData(x, n, 1);
dmat->Info().feature_types.HostVector().push_back(FeatureType::kCategorical);
if (weighted) {
std::vector<float> weights(n, 0);
SimpleLCG lcg;
SimpleRealUniformDistribution<float> dist(0, 1);
for (auto& v : weights) {
v = dist(&lcg);
}
dmat->Info().weights_.HostVector() = weights;
}
ASSERT_EQ(dmat->Info().feature_types.Size(), 1);
auto cuts = DeviceSketch(0, dmat.get(), num_bins);
std::sort(x.begin(), x.end());
@@ -146,7 +157,8 @@ void TestCategoricalSketch(size_t n, size_t num_categories, int32_t num_bins) {
}
TEST(HistUtil, DeviceSketchCategoricalFeatures) {
TestCategoricalSketch(1000, 256, 32);
TestCategoricalSketch(1000, 256, 32, false);
TestCategoricalSketch(1000, 256, 32, true);
}
TEST(HistUtil, DeviceSketchMultipleColumns) {