Implement GPU predict leaf. (#6187)

This commit is contained in:
Jiaming Yuan
2020-11-11 17:33:47 +08:00
committed by GitHub
parent 7f101d1b33
commit 8a17610666
12 changed files with 252 additions and 42 deletions

View File

@@ -345,7 +345,7 @@ class CPUPredictor : public Predictor {
}
}
void PredictLeaf(DMatrix* p_fmat, std::vector<bst_float>* out_preds,
void PredictLeaf(DMatrix* p_fmat, HostDeviceVector<bst_float>* out_preds,
const gbm::GBTreeModel& model, unsigned ntree_limit) override {
const int nthread = omp_get_max_threads();
InitThreadTemp(nthread, model.learner_model_param->num_feature, &this->thread_temp_);
@@ -355,7 +355,7 @@ class CPUPredictor : public Predictor {
if (ntree_limit == 0 || ntree_limit > model.trees.size()) {
ntree_limit = static_cast<unsigned>(model.trees.size());
}
std::vector<bst_float>& preds = *out_preds;
std::vector<bst_float>& preds = out_preds->HostVector();
preds.resize(info.num_row_ * ntree_limit);
// start collecting the prediction
for (const auto &batch : p_fmat->GetBatches<SparsePage>()) {