Skip to content

Commit

Permalink
Move prediction cache to Learner. (#5220)
Browse files Browse the repository at this point in the history 
* Move prediction cache into Learner.

* Clean-ups

- Remove duplicated cache in Learner and GBM.
- Remove ad-hoc fix of invalid cache.
- Remove `PredictFromCache` in predictors.
- Remove prediction cache for linear altogether, as it's only moving the
  prediction into training process but doesn't provide any actual overall speed
  gain.
- The cache is now unique to Learner, which means the ownership is no longer
  shared by any other components.

* Changes

- Add version to prediction cache.
- Use weak ptr to check expired DMatrix.
- Pass shared pointer instead of raw pointer.
  • Loading branch information
@trivialfis
trivialfis authored Feb 14, 2020
1 parent 24ad9de commit c35cdec
Show file tree
Hide file tree
Showing 19 changed files with 454 additions and 369 deletions.
28 changes: 16 additions & 12 deletions include/xgboost/gbm.h
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*!
* Copyright by Contributors
* Copyright 2014-2020 by Contributors
* \file gbm.h
* \brief Interface of gradient booster,
* that learns through gradient statistics.
Expand All @@ -18,6 +18,7 @@
#include <utility>
#include <string>
#include <functional>
#include <unordered_map>
#include <memory>

namespace xgboost {
Expand All @@ -28,6 +29,8 @@ class ObjFunction;

struct GenericParameter;
struct LearnerModelParam;
struct PredictionCacheEntry;
class PredictionContainer;

/*!
* \brief interface of gradient boosting model.
Expand All @@ -38,7 +41,7 @@ class GradientBooster : public Model, public Configurable {

public:
/*! \brief virtual destructor */
virtual ~GradientBooster() = default;
~GradientBooster() override = default;
/*!
* \brief Set the configuration of gradient boosting.
* User must call configure once before InitModel and Training.
Expand Down Expand Up @@ -71,19 +74,22 @@ class GradientBooster : public Model, public Configurable {
* \param obj The objective function, optional, can be nullptr when use customized version
* the booster may change content of gpair
*/
virtual void DoBoost(DMatrix* p_fmat,
HostDeviceVector<GradientPair>* in_gpair,
ObjFunction* obj = nullptr) = 0;
virtual void DoBoost(DMatrix* p_fmat, HostDeviceVector<GradientPair>* in_gpair,
PredictionCacheEntry *prediction) = 0;

/*!
* \brief generate predictions for given feature matrix
* \param dmat feature matrix
* \param out_preds output vector to hold the predictions
* \param ntree_limit limit the number of trees used in prediction, when it equals 0, this means
* we do not limit number of trees, this parameter is only valid for gbtree, but not for gblinear
* \param training Whether the prediction value is used for training. For dart booster
* drop out is performed during training.
* \param ntree_limit limit the number of trees used in prediction,
* when it equals 0, this means we do not limit
* number of trees, this parameter is only valid
* for gbtree, but not for gblinear
*/
virtual void PredictBatch(DMatrix* dmat,
HostDeviceVector<bst_float>* out_preds,
PredictionCacheEntry* out_preds,
bool training,
unsigned ntree_limit = 0) = 0;
/*!
Expand Down Expand Up @@ -158,8 +164,7 @@ class GradientBooster : public Model, public Configurable {
static GradientBooster* Create(
const std::string& name,
GenericParameter const* generic_param,
LearnerModelParam const* learner_model_param,
const std::vector<std::shared_ptr<DMatrix> >& cache_mats);
LearnerModelParam const* learner_model_param);

static void AssertGPUSupport() {
#ifndef XGBOOST_USE_CUDA
Expand All @@ -174,8 +179,7 @@ class GradientBooster : public Model, public Configurable {
struct GradientBoosterReg
: public dmlc::FunctionRegEntryBase<
GradientBoosterReg,
std::function<GradientBooster* (const std::vector<std::shared_ptr<DMatrix> > &cached_mats,
LearnerModelParam const* learner_model_param)> > {
std::function<GradientBooster* (LearnerModelParam const* learner_model_param)> > {
};

/*!
Expand Down
108 changes: 71 additions & 37 deletions include/xgboost/predictor.h
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*!
* Copyright by Contributors
* Copyright 2017-2020 by Contributors
* \file predictor.h
* \brief Interface of predictor,
* performs predictions for a gradient booster.
Expand Down Expand Up @@ -32,47 +32,83 @@ namespace xgboost {
* \brief Contains pointer to input matrix and associated cached predictions.
*/
struct PredictionCacheEntry {
std::shared_ptr<DMatrix> data;
// A storage for caching prediction values
HostDeviceVector<bst_float> predictions;
// The version of current cache, corresponding number of layers of trees
uint32_t version;
// A weak pointer for checking whether the DMatrix object has expired.
std::weak_ptr< DMatrix > ref;

PredictionCacheEntry() : version { 0 } {}
/* \brief Update the cache entry by number of versions.
*
* \param v Added versions.
*/
void Update(uint32_t v) {
version += v;
}
};

/* \brief A container for managed prediction caches.
*/
class PredictionContainer {
std::unordered_map<DMatrix *, PredictionCacheEntry> container_;
void ClearExpiredEntries();

public:
PredictionContainer() = default;
/* \brief Add a new DMatrix to the cache, at the same time this function will clear out
* all expired caches by checking the `std::weak_ptr`. Caching an existing
* DMatrix won't renew it.
*
* Passing in a `shared_ptr` is critical here. First to create a `weak_ptr` inside the
* entry this shared pointer is necessary. More importantly, the life time of this
* cache is tied to the shared pointer.
*
* Another way to make a safe cache is create a proxy to this entry, with anther shared
* pointer defined inside, and pass this proxy around instead of the real entry. But
* seems to be too messy. In XGBoost, functions like `UpdateOneIter` will have
* (memory) safe access to the DMatrix as long as it's passed in as a `shared_ptr`.
*
* \param m shared pointer to the DMatrix that needs to be cached.
* \param device Which device should the cache be allocated on. Pass
* GenericParameter::kCpuId for CPU or positive integer for GPU id.
*
* \return the cache entry for passed in DMatrix, either an existing cache or newly
* created.
*/
PredictionCacheEntry& Cache(std::shared_ptr<DMatrix> m, int32_t device);
/* \brief Get a prediction cache entry. This entry must be already allocated by `Cache`
* method. Otherwise a dmlc::Error is thrown.
*
* \param m pointer to the DMatrix.
* \return The prediction cache for passed in DMatrix.
*/
PredictionCacheEntry& Entry(DMatrix* m);
/* \brief Get a const reference to the underlying hash map. Clear expired caches before
* returning.
*/
decltype(container_) const& Container();
};

/**
* \class Predictor
*
* \brief Performs prediction on individual training instances or batches of
* instances for GBTree. The predictor also manages a prediction cache
* associated with input matrices. If possible, it will use previously
* calculated predictions instead of calculating new predictions.
* Prediction functions all take a GBTreeModel and a DMatrix as input and
* output a vector of predictions. The predictor does not modify any state of
* the model itself.
* \brief Performs prediction on individual training instances or batches of instances for
* GBTree. Prediction functions all take a GBTreeModel and a DMatrix as input and
* output a vector of predictions. The predictor does not modify any state of the
* model itself.
*/

class Predictor {
protected:
/*
* \brief Runtime parameters.
*/
GenericParameter const* generic_param_;
/**
* \brief Map of matrices and associated cached predictions to facilitate
* storing and looking up predictions.
*/
std::shared_ptr<std::unordered_map<DMatrix*, PredictionCacheEntry>> cache_;

std::unordered_map<DMatrix*, PredictionCacheEntry>::iterator FindCache(DMatrix const* dmat) {
auto cache_emtry = std::find_if(
cache_->begin(), cache_->end(),
[dmat](std::pair<DMatrix *, PredictionCacheEntry const &> const &kv) {
return kv.second.data.get() == dmat;
});
return cache_emtry;
}

public:
Predictor(GenericParameter const* generic_param,
std::shared_ptr<std::unordered_map<DMatrix*, PredictionCacheEntry>> cache) :
generic_param_{generic_param}, cache_{cache} {}
explicit Predictor(GenericParameter const* generic_param) :
generic_param_{generic_param} {}
virtual ~Predictor() = default;

/**
Expand All @@ -91,12 +127,11 @@ class Predictor {
* \param model The model to predict from.
* \param tree_begin The tree begin index.
* \param ntree_limit (Optional) The ntree limit. 0 means do not
* limit trees.
* limit trees.
*/

virtual void PredictBatch(DMatrix* dmat, HostDeviceVector<bst_float>* out_preds,
virtual void PredictBatch(DMatrix* dmat, PredictionCacheEntry* out_preds,
const gbm::GBTreeModel& model, int tree_begin,
unsigned ntree_limit = 0) = 0;
uint32_t const ntree_limit = 0) = 0;

/**
* \fn virtual void Predictor::UpdatePredictionCache( const gbm::GBTreeModel
Expand All @@ -116,7 +151,9 @@ class Predictor {
virtual void UpdatePredictionCache(
const gbm::GBTreeModel& model,
std::vector<std::unique_ptr<TreeUpdater>>* updaters,
int num_new_trees) = 0;
int num_new_trees,
DMatrix* m,
PredictionCacheEntry* predts) = 0;

/**
* \fn virtual void Predictor::PredictInstance( const SparsePage::Inst&
Expand Down Expand Up @@ -200,18 +237,15 @@ class Predictor {
* \param cache Pointer to prediction cache.
*/
static Predictor* Create(
std::string const& name, GenericParameter const* generic_param,
std::shared_ptr<std::unordered_map<DMatrix*, PredictionCacheEntry>> cache);
std::string const& name, GenericParameter const* generic_param);
};

/*!
* \brief Registry entry for predictor.
*/
struct PredictorReg
: public dmlc::FunctionRegEntryBase<
PredictorReg, std::function<Predictor*(
GenericParameter const*,
std::shared_ptr<std::unordered_map<DMatrix*, PredictionCacheEntry>>)>> {};
PredictorReg, std::function<Predictor*(GenericParameter const*)>> {};

#define XGBOOST_REGISTER_PREDICTOR(UniqueId, Name) \
static DMLC_ATTRIBUTE_UNUSED ::xgboost::PredictorReg& \
Expand Down
2 changes: 1 addition & 1 deletion include/xgboost/tree_model.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -158,7 +158,7 @@ class RegTree : public Model {
}
/*! \brief whether this node is deleted */
XGBOOST_DEVICE bool IsDeleted() const {
return sindex_ == std::numeric_limits<unsigned>::max();
return sindex_ == std::numeric_limits<uint32_t>::max();
}
/*! \brief whether current node is root */
XGBOOST_DEVICE bool IsRoot() const { return parent_ == kInvalidNodeId; }
Expand Down
66 changes: 12 additions & 54 deletions src/gbm/gblinear.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@

#include "xgboost/gbm.h"
#include "xgboost/json.h"
#include "xgboost/predictor.h"
#include "xgboost/linear_updater.h"
#include "xgboost/logging.h"
#include "xgboost/learner.h"
Expand Down Expand Up @@ -50,21 +51,14 @@ struct GBLinearTrainParam : public XGBoostParameter<GBLinearTrainParam> {
*/
class GBLinear : public GradientBooster {
public:
explicit GBLinear(const std::vector<std::shared_ptr<DMatrix> > &cache,
LearnerModelParam const* learner_model_param)
explicit GBLinear(LearnerModelParam const* learner_model_param)
: learner_model_param_{learner_model_param},
model_{learner_model_param_},
previous_model_{learner_model_param_},
sum_instance_weight_(0),
sum_weight_complete_(false),
is_converged_(false) {
// Add matrices to the prediction cache
for (auto &d : cache) {
PredictionCacheEntry e;
e.data = d;
cache_[d.get()] = std::move(e);
}
}
is_converged_(false) {}

void Configure(const Args& cfg) override {
if (model_.weight.size() == 0) {
model_.Configure(cfg);
Expand Down Expand Up @@ -118,7 +112,7 @@ class GBLinear : public GradientBooster {

void DoBoost(DMatrix *p_fmat,
HostDeviceVector<GradientPair> *in_gpair,
ObjFunction* obj) override {
PredictionCacheEntry* predt) override {
monitor_.Start("DoBoost");

model_.LazyInitModel();
Expand All @@ -127,28 +121,19 @@ class GBLinear : public GradientBooster {
if (!this->CheckConvergence()) {
updater_->Update(in_gpair, p_fmat, &model_, sum_instance_weight_);
}
this->UpdatePredictionCache();

monitor_.Stop("DoBoost");
}

void PredictBatch(DMatrix *p_fmat,
HostDeviceVector<bst_float> *out_preds,
PredictionCacheEntry *predts,
bool training,
unsigned ntree_limit) override {
monitor_.Start("PredictBatch");
auto* out_preds = &predts->predictions;
CHECK_EQ(ntree_limit, 0U)
<< "GBLinear::Predict ntrees is only valid for gbtree predictor";

// Try to predict from cache
auto it = cache_.find(p_fmat);
if (it != cache_.end() && it->second.predictions.size() != 0) {
std::vector<bst_float> &y = it->second.predictions;
out_preds->Resize(y.size());
std::copy(y.begin(), y.end(), out_preds->HostVector().begin());
} else {
this->PredictBatchInternal(p_fmat, &out_preds->HostVector());
}
this->PredictBatchInternal(p_fmat, &out_preds->HostVector());
monitor_.Stop("PredictBatch");
}
// add base margin
Expand Down Expand Up @@ -258,25 +243,15 @@ class GBLinear : public GradientBooster {
const size_t ridx = batch.base_rowid + i;
// loop over output groups
for (int gid = 0; gid < ngroup; ++gid) {
bst_float margin = (base_margin.size() != 0) ?
bst_float margin =
(base_margin.size() != 0) ?
base_margin[ridx * ngroup + gid] : learner_model_param_->base_score;
this->Pred(batch[i], &preds[ridx * ngroup], gid, margin);
}
}
}
monitor_.Stop("PredictBatchInternal");
}
void UpdatePredictionCache() {
// update cache entry
for (auto &kv : cache_) {
PredictionCacheEntry &e = kv.second;
if (e.predictions.size() == 0) {
size_t n = model_.learner_model_param_->num_output_group * e.data->Info().num_row_;
e.predictions.resize(n);
}
this->PredictBatchInternal(e.data.get(), &e.predictions);
}
}

bool CheckConvergence() {
if (param_.tolerance == 0.0f) return false;
Expand Down Expand Up @@ -327,32 +302,15 @@ class GBLinear : public GradientBooster {
bool sum_weight_complete_;
common::Monitor monitor_;
bool is_converged_;

/**
* \struct PredictionCacheEntry
*
* \brief Contains pointer to input matrix and associated cached predictions.
*/
struct PredictionCacheEntry {
std::shared_ptr<DMatrix> data;
std::vector<bst_float> predictions;
};

/**
* \brief Map of matrices and associated cached predictions to facilitate
* storing and looking up predictions.
*/
std::unordered_map<DMatrix*, PredictionCacheEntry> cache_;
};

// register the objective functions
DMLC_REGISTER_PARAMETER(GBLinearTrainParam);

XGBOOST_REGISTER_GBM(GBLinear, "gblinear")
.describe("Linear booster, implement generalized linear model.")
.set_body([](const std::vector<std::shared_ptr<DMatrix> > &cache,
LearnerModelParam const* booster_config) {
return new GBLinear(cache, booster_config);
.set_body([](LearnerModelParam const* booster_config) {
return new GBLinear(booster_config);
});
} // namespace gbm
} // namespace xgboost
Loading

0 comments on commit c35cdec

Please sign in to comment.