Support hessian in host sketch container. (#7081)

Prepare for migrating approx onto hist's codebase.
This commit is contained in:
Jiaming Yuan
2021-07-08 16:33:58 +08:00
committed by GitHub
parent 84d359efb8
commit 77f6cf2d13
8 changed files with 238 additions and 53 deletions

View File

@@ -110,7 +110,8 @@ class HistogramCuts {
}
};
inline HistogramCuts SketchOnDMatrix(DMatrix *m, int32_t max_bins) {
inline HistogramCuts SketchOnDMatrix(DMatrix *m, int32_t max_bins,
std::vector<float> const &hessian = {}) {
HistogramCuts out;
auto const& info = m->Info();
const auto threads = omp_get_max_threads();
@@ -127,9 +128,9 @@ inline HistogramCuts SketchOnDMatrix(DMatrix *m, int32_t max_bins) {
}
}
HostSketchContainer container(reduced, max_bins,
HostSketchContainer::UseGroup(info));
HostSketchContainer::UseGroup(info), threads);
for (auto const &page : m->GetBatches<SparsePage>()) {
container.PushRowPage(page, info);
container.PushRowPage(page, info, hessian);
}
container.MakeCuts(&out);
return out;