Implement a full blown cache with validation. Add failed test.

dmlc · Jan 22, 2020 · 04d4dc5 · 04d4dc5
1 parent 47c4712
commit 04d4dc5
Show file tree

Hide file tree

Showing 16 changed files with 112 additions and 80 deletions.
diff --git a/include/xgboost/gbm.h b/include/xgboost/gbm.h
@@ -30,6 +30,7 @@ class ObjFunction;
 struct GenericParameter;
 struct LearnerModelParam;
 struct PredictionCacheEntry;
+struct PredictionContainer;
 
 /*!
  * \brief interface of gradient boosting model.
@@ -161,7 +162,7 @@ class GradientBooster : public Model, public Configurable {
       const std::string& name,
       GenericParameter const* generic_param,
       LearnerModelParam const* learner_model_param,
-      std::shared_ptr<std::unordered_map<DMatrix*, PredictionCacheEntry>> cache);
+      std::shared_ptr< PredictionContainer > cache);
 
   static void AssertGPUSupport() {
 #ifndef XGBOOST_USE_CUDA
@@ -177,7 +178,7 @@ struct GradientBoosterReg
     : public dmlc::FunctionRegEntryBase<
   GradientBoosterReg,
   std::function<GradientBooster* (
-      std::shared_ptr<std::unordered_map<DMatrix*, PredictionCacheEntry>> cache,
+      std::shared_ptr< PredictionContainer > cache,
       LearnerModelParam const* learner_model_param)> > {
 };
 

diff --git a/include/xgboost/predictor.h b/include/xgboost/predictor.h
@@ -32,8 +32,43 @@ namespace xgboost {
  * \brief Contains pointer to input matrix and associated cached predictions.
  */
 struct PredictionCacheEntry {
-  std::shared_ptr<DMatrix> data;
+  bool valid;
   HostDeviceVector<bst_float> predictions;
+
+  PredictionCacheEntry() : valid { false } {}
+};
+
+struct PredictionContainer {
+  std::unordered_map<DMatrix*, PredictionCacheEntry> container;
+
+ public:
+  PredictionContainer() = default;
+
+  HostDeviceVector<float> const*  FetchConst(DMatrix* m) const {
+    if (container.find(m) != container.cend()) {
+      return &(container.at(m).predictions);
+    }
+    return nullptr;
+  }
+  HostDeviceVector<float>* Fetch(DMatrix* m) {
+    if (container.find(m) != container.cend()) {
+      return &(container.at(m).predictions);
+    }
+    container[m].valid = false;
+    return nullptr;
+  }
+  void SetValid(DMatrix* m) {
+    container.at(m).valid = true;
+  }
+  void Obsolate() {
+    for (auto& d : container) {
+      d.second.valid = false;
+    }
+  }
+
+  decltype(container) const& Container() {
+    return container;
+  }
 };
 
 /**
@@ -58,21 +93,12 @@ class Predictor {
    * \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;
-  }
+  std::shared_ptr< PredictionContainer > cache_;
 
  public:
   Predictor(GenericParameter const* generic_param,
-            std::shared_ptr<std::unordered_map<DMatrix*, PredictionCacheEntry>> cache) :
-      generic_param_{generic_param}, cache_{cache} {}
+            std::shared_ptr< PredictionContainer > cache) :
+      generic_param_{generic_param}, cache_{std::move(cache)} {}
   virtual ~Predictor() = default;
 
   /**
@@ -201,7 +227,7 @@ class Predictor {
    */
   static Predictor* Create(
       std::string const& name, GenericParameter const* generic_param,
-      std::shared_ptr<std::unordered_map<DMatrix*, PredictionCacheEntry>> cache);
+      std::shared_ptr< PredictionContainer > cache);
 };
 
 /*!
@@ -211,7 +237,7 @@ struct PredictorReg
     : public dmlc::FunctionRegEntryBase<
   PredictorReg, std::function<Predictor*(
       GenericParameter const*,
-      std::shared_ptr<std::unordered_map<DMatrix*, PredictionCacheEntry>>)>> {};
+      std::shared_ptr< PredictionContainer >)>> {};
 
 #define XGBOOST_REGISTER_PREDICTOR(UniqueId, Name)      \
   static DMLC_ATTRIBUTE_UNUSED ::xgboost::PredictorReg& \

diff --git a/src/common/observer.h b/src/common/observer.h
@@ -24,7 +24,7 @@ namespace xgboost {
  *  similiar facilities.
  */
 class TrainingObserver {
-#if !defined(XGBOOST_USE_DEBUG_OUTPUT)
+#if defined(XGBOOST_USE_DEBUG_OUTPUT)
   bool constexpr static observe_ {true};
 #else
   bool constexpr static observe_ {false};

diff --git a/src/gbm/gblinear.cc b/src/gbm/gblinear.cc
@@ -51,14 +51,15 @@ struct GBLinearTrainParam : public XGBoostParameter<GBLinearTrainParam> {
  */
 class GBLinear : public GradientBooster {
  public:
-  explicit GBLinear(std::shared_ptr<std::unordered_map<DMatrix*, PredictionCacheEntry>> const cache,
+  explicit GBLinear(std::shared_ptr< PredictionContainer > const cache,
                     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) {}
+        is_converged_(false),
+        cache_{cache} {}
 
   void Configure(const Args& cfg) override {
     if (model_.weight.size() == 0) {
@@ -129,9 +130,9 @@ class GBLinear : public GradientBooster {
         << "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.HostVector();
+    auto it = cache_->Container().find(p_fmat);
+    if (it != cache_->Container().cend() && it->second.predictions.Size() != 0) {
+      std::vector<bst_float> &y = cache_->Fetch(it->first)->HostVector();
       out_preds->Resize(y.size());
       std::copy(y.begin(), y.end(), out_preds->HostVector().begin());
     } else {
@@ -253,13 +254,15 @@ class GBLinear : public GradientBooster {
   }
   void UpdatePredictionCache() {
     // update cache entry
-    for (auto &kv : cache_) {
-      PredictionCacheEntry &e = kv.second;
+    for (auto const &kv : cache_->Container()) {
+      PredictionCacheEntry const &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);
+        size_t n = model_.learner_model_param_->num_output_group * kv.first->Info().num_row_;
+        auto e = cache_->Fetch(kv.first);
+        e->Resize(n);
       }
-      this->PredictBatchInternal(e.data.get(), &(e.predictions.HostVector()));
+      auto prediction = cache_->Fetch(kv.first);
+      this->PredictBatchInternal(kv.first, &prediction->HostVector());
     }
   }
 
@@ -317,15 +320,15 @@ class GBLinear : public GradientBooster {
    * \brief Map of matrices and associated cached predictions to facilitate
    * storing and looking up predictions.
    */
-  std::unordered_map<DMatrix*, PredictionCacheEntry> cache_;
+  std::shared_ptr< PredictionContainer > cache_;
 };
 
 // register the objective functions
 DMLC_REGISTER_PARAMETER(GBLinearTrainParam);
 
 XGBOOST_REGISTER_GBM(GBLinear, "gblinear")
     .describe("Linear booster, implement generalized linear model.")
-    .set_body([](std::shared_ptr<std::unordered_map<DMatrix*, PredictionCacheEntry>> cache,
+    .set_body([](std::shared_ptr< PredictionContainer > cache,
                  LearnerModelParam const* booster_config) {
       return new GBLinear(cache, booster_config);
     });

diff --git a/src/gbm/gbm.cc b/src/gbm/gbm.cc
@@ -21,7 +21,7 @@ GradientBooster* GradientBooster::Create(
     const std::string& name,
     GenericParameter const* generic_param,
     LearnerModelParam const* learner_model_param,
-    std::shared_ptr<std::unordered_map<DMatrix*, PredictionCacheEntry>> cache) {
+    std::shared_ptr< PredictionContainer > cache) {
   auto *e = ::dmlc::Registry< ::xgboost::GradientBoosterReg>::Get()->Find(name);
   if (e == nullptr) {
     LOG(FATAL) << "Unknown gbm type " << name;

diff --git a/src/gbm/gbtree.cc b/src/gbm/gbtree.cc
@@ -354,7 +354,7 @@ void GBTree::SaveModel(Json* p_out) const {
 
 class Dart : public GBTree {
  public:
-  explicit Dart(std::shared_ptr<std::unordered_map<DMatrix*, PredictionCacheEntry>> cache,
+  explicit Dart(std::shared_ptr< PredictionContainer > cache,
                 LearnerModelParam const* booster_config) :
       GBTree(cache, booster_config) {}
 
@@ -683,14 +683,14 @@ DMLC_REGISTER_PARAMETER(DartTrainParam);
 
 XGBOOST_REGISTER_GBM(GBTree, "gbtree")
 .describe("Tree booster, gradient boosted trees.")
-.set_body([](std::shared_ptr<std::unordered_map<DMatrix*, PredictionCacheEntry>> cache,
+.set_body([](std::shared_ptr< PredictionContainer > cache,
              LearnerModelParam const* booster_config) {
     auto* p = new GBTree(cache, booster_config);
     return p;
   });
 XGBOOST_REGISTER_GBM(Dart, "dart")
 .describe("Tree booster, dart.")
-.set_body([](std::shared_ptr<std::unordered_map<DMatrix*, PredictionCacheEntry>> cache,
+.set_body([](std::shared_ptr< PredictionContainer > cache,
              LearnerModelParam const* booster_config) {
     GBTree* p = new Dart(cache, booster_config);
     return p;

diff --git a/src/gbm/gbtree.h b/src/gbm/gbtree.h
@@ -151,7 +151,7 @@ struct DartTrainParam : public XGBoostParameter<DartTrainParam> {
 // gradient boosted trees
 class GBTree : public GradientBooster {
  public:
-  explicit GBTree(std::shared_ptr<std::unordered_map<DMatrix*, PredictionCacheEntry>> cache,
+  explicit GBTree(std::shared_ptr< PredictionContainer > cache,
                   LearnerModelParam const* booster_config) :
       cache_{std::move(cache)}, model_(booster_config) {}
 
@@ -320,7 +320,7 @@ class GBTree : public GradientBooster {
    * \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::shared_ptr< PredictionContainer > cache_;
   GBTreeModel model_;
   // training parameter
   GBTreeTrainParam tparam_;

diff --git a/src/learner.cc b/src/learner.cc
@@ -188,12 +188,12 @@ void GenericParameter::ConfigureGpuId(bool require_gpu) {
  */
 class LearnerImpl : public Learner {
  public:
-  explicit LearnerImpl(std::vector<std::shared_ptr<DMatrix> >  cache)
+  explicit LearnerImpl(std::vector<std::shared_ptr<DMatrix> > cache)
       : need_configuration_{true} {
     monitor_.Init("Learner");
-    cache_ = std::make_shared<std::unordered_map<DMatrix*, PredictionCacheEntry>>();
+    cache_ = std::make_shared<PredictionContainer>();
     for (std::shared_ptr<DMatrix> const& d : cache) {
-      (*cache_)[d.get()].data = d;
+      cache_->Fetch(d.get());
     }
   }
   // Configuration before data is known.
@@ -698,15 +698,15 @@ class LearnerImpl : public Learner {
     this->CheckDataSplitMode();
     this->ValidateDMatrix(train);
 
-    auto& predt = (*cache_)[train].predictions;
+    auto* predt = cache_->Fetch(train);
 
     monitor_.Start("PredictRaw");
-    this->PredictRaw(train, &predt, true);
+    this->PredictRaw(train, predt, true);
     monitor_.Stop("PredictRaw");
     TrainingObserver::Instance().Observe(predt, "Predictions");
 
     monitor_.Start("GetGradient");
-    obj_->GetGradient(predt, train->Info(), iter, &gpair_);
+    obj_->GetGradient(*predt, train->Info(), iter, &gpair_);
     monitor_.Stop("GetGradient");
     TrainingObserver::Instance().Observe(gpair_, "Gradients");
 
@@ -741,13 +741,13 @@ class LearnerImpl : public Learner {
     }
     for (size_t i = 0; i < data_sets.size(); ++i) {
       DMatrix * dmat = data_sets[i];
-      auto& predt = (*cache_)[dmat].predictions;
+      auto* predt = cache_->Fetch(dmat);
       this->ValidateDMatrix(dmat);
-      this->PredictRaw(data_sets[i], &predt, false);
-      obj_->EvalTransform(&predt);
+      this->PredictRaw(data_sets[i], predt, false);
+      obj_->EvalTransform(predt);
       for (auto& ev : metrics_) {
         os << '\t' << data_names[i] << '-' << ev->Name() << ':'
-           << ev->Eval(predt, data_sets[i]->Info(),
+           << ev->Eval(*predt, data_sets[i]->Info(),
                        tparam_.dsplit == DataSplitMode::kRow);
       }
     }
@@ -846,8 +846,7 @@ class LearnerImpl : public Learner {
   void PredictRaw(DMatrix* data, HostDeviceVector<bst_float>* out_preds,
                   bool training,
                   unsigned ntree_limit = 0) const {
-    CHECK(gbm_ != nullptr)
-        << "Predict must happen after Load or configuration";
+    CHECK(gbm_ != nullptr) << "Predict must happen after Load or configuration";
     this->ValidateDMatrix(data);
     gbm_->PredictBatch(data, out_preds, training, ntree_limit);
   }
@@ -904,7 +903,7 @@ class LearnerImpl : public Learner {
     // estimate feature bound
     // TODO(hcho3): Change num_feature to 64-bit integer
     unsigned num_feature = 0;
-    for (auto & matrix : *cache_) {
+    for (auto & matrix : cache_->Container()) {
       CHECK(matrix.first != nullptr);
       const uint64_t num_col = matrix.first->Info().num_col_;
       CHECK_LE(num_col, static_cast<uint64_t>(std::numeric_limits<unsigned>::max()))
@@ -979,7 +978,7 @@ class LearnerImpl : public Learner {
   /*! \brief random number transformation seed. */
   static int32_t constexpr kRandSeedMagic = 127;
   // internal cached dmatrix for prediction.
-  std::shared_ptr<std::unordered_map<DMatrix*, PredictionCacheEntry>> cache_;
+  std::shared_ptr< PredictionContainer > cache_;
 
   common::Monitor monitor_;
 

diff --git a/src/predictor/cpu_predictor.cc b/src/predictor/cpu_predictor.cc
@@ -109,9 +109,10 @@ class CPUPredictor : public Predictor {
     CHECK(cache_);
     if (ntree_limit == 0 ||
         ntree_limit * model.learner_model_param_->num_output_group >= model.trees.size()) {
-      auto it = cache_->find(dmat);
-      if (it != cache_->end()) {
+      auto it = cache_->Container().find(dmat);
+      if (it != cache_->Container().end()) {
         const HostDeviceVector<bst_float>& y = it->second.predictions;
+        CHECK(cache_->Container().at(dmat).valid);
         if (y.Size() != 0) {
           out_preds->Resize(y.Size());
           std::copy(y.HostVector().begin(), y.HostVector().end(),
@@ -157,7 +158,7 @@ class CPUPredictor : public Predictor {
 
  public:
   CPUPredictor(GenericParameter const* generic_param,
-               std::shared_ptr<std::unordered_map<DMatrix*, PredictionCacheEntry>> cache) :
+               std::shared_ptr< PredictionContainer > cache) :
       Predictor::Predictor{generic_param, cache} {}
   void PredictBatch(DMatrix* dmat, HostDeviceVector<bst_float>* out_preds,
                     const gbm::GBTreeModel& model, int tree_begin,
@@ -183,22 +184,22 @@ class CPUPredictor : public Predictor {
     int old_ntree = model.trees.size() - num_new_trees;
     // update cache entry
     CHECK(cache_);
-    for (auto& kv : (*cache_)) {
-      PredictionCacheEntry& e = kv.second;
-
+    for (auto& kv : cache_->Container()) {
+      PredictionCacheEntry const& e = kv.second;
+      auto* out = cache_->Fetch(kv.first);
       if (e.predictions.Size() == 0) {
-        InitOutPredictions(e.data->Info(), &(e.predictions), model);
-        PredLoopInternal(e.data.get(), &(e.predictions.HostVector()), model, 0,
+        InitOutPredictions(kv.first->Info(), out, model);
+        PredLoopInternal(kv.first, &(out->HostVector()), model, 0,
                          model.trees.size());
       } else if (model.learner_model_param_->num_output_group == 1 && updaters->size() > 0 &&
                  num_new_trees == 1 &&
-                 updaters->back()->UpdatePredictionCache(e.data.get(),
-                                                         &(e.predictions))) {
+                 updaters->back()->UpdatePredictionCache(kv.first, out)) {
         {}  // do nothing
       } else {
-        PredLoopInternal(e.data.get(), &(e.predictions.HostVector()), model, old_ntree,
+        PredLoopInternal(kv.first, &(out->HostVector()), model, old_ntree,
                          model.trees.size());
       }
+      cache_->SetValid(kv.first);
     }
   }
 
@@ -378,7 +379,7 @@ class CPUPredictor : public Predictor {
 XGBOOST_REGISTER_PREDICTOR(CPUPredictor, "cpu_predictor")
 .describe("Make predictions using CPU.")
 .set_body([](GenericParameter const* generic_param,
-             std::shared_ptr<std::unordered_map<DMatrix*, PredictionCacheEntry>> cache) {
+             std::shared_ptr< PredictionContainer > cache) {
             return new CPUPredictor(generic_param, cache);
           });
 }  // namespace predictor