predictor.h

Go to the documentation of this file.

 
#pragma once
#include <xgboost/base.h>
#include <xgboost/cache.h>    // for DMatrixCache
#include <xgboost/context.h>  // for Context
#include <xgboost/context.h>
#include <xgboost/data.h>
#include <xgboost/host_device_vector.h>
 
#include <functional>  // for function
#include <memory>      // for shared_ptr
#include <string>
#include <utility>  // for make_pair
#include <vector>
 
// Forward declarations
namespace xgboost::gbm {
struct GBTreeModel;
}  // namespace xgboost::gbm
 
namespace xgboost {
struct PredictionCacheEntry {
  // A storage for caching prediction values
  HostDeviceVector<bst_float> predictions;
  // The version of current cache, corresponding number of layers of trees
  std::uint32_t version{0};
 
  PredictionCacheEntry() = default;
  void Update(std::uint32_t v) { version += v; }
  void Reset() { version = 0; }
};
 
class PredictionContainer : public DMatrixCache<PredictionCacheEntry> {
  // We cache up to 64 DMatrix for all threads
  std::size_t static constexpr DefaultSize() { return 64; }
 
 public:
  PredictionContainer() : DMatrixCache<PredictionCacheEntry>{DefaultSize()} {}
  PredictionCacheEntry& Cache(std::shared_ptr<DMatrix> m, std::int32_t device) {
    auto p_cache = this->CacheItem(m);
    if (device != Context::kCpuId) {
      p_cache->predictions.SetDevice(device);
    }
    return *p_cache;
  }
};
 
class Predictor {
 protected:
  Context const* ctx_;
 
 public:
  explicit Predictor(Context const* ctx) : ctx_{ctx} {}
 
  virtual ~Predictor() = default;
 
  virtual void Configure(Args const&);
 
  void InitOutPredictions(const MetaInfo& info, HostDeviceVector<bst_float>* out_predt,
                          const gbm::GBTreeModel& model) const;
 
  virtual void PredictBatch(DMatrix* dmat, PredictionCacheEntry* out_preds,
                            const gbm::GBTreeModel& model, uint32_t tree_begin,
                            uint32_t tree_end = 0) const = 0;
 
  virtual bool InplacePredict(std::shared_ptr<DMatrix> p_fmat, const gbm::GBTreeModel& model,
                              float missing, PredictionCacheEntry* out_preds,
                              uint32_t tree_begin = 0, uint32_t tree_end = 0) const = 0;
  virtual void PredictInstance(const SparsePage::Inst& inst,
                               std::vector<bst_float>* out_preds,
                               const gbm::GBTreeModel& model,
                               unsigned tree_end = 0,
                               bool is_column_split = false) const = 0;
 
  virtual void PredictLeaf(DMatrix* dmat, HostDeviceVector<bst_float>* out_preds,
                           const gbm::GBTreeModel& model,
                           unsigned tree_end = 0) const = 0;
 
  virtual void
  PredictContribution(DMatrix *dmat, HostDeviceVector<bst_float> *out_contribs,
                      const gbm::GBTreeModel &model, unsigned tree_end = 0,
                      std::vector<bst_float> const *tree_weights = nullptr,
                      bool approximate = false, int condition = 0,
                      unsigned condition_feature = 0) const = 0;
 
  virtual void PredictInteractionContributions(
      DMatrix *dmat, HostDeviceVector<bst_float> *out_contribs,
      const gbm::GBTreeModel &model, unsigned tree_end = 0,
      std::vector<bst_float> const *tree_weights = nullptr,
      bool approximate = false) const = 0;
 
  static Predictor* Create(std::string const& name, Context const* ctx);
};
 
struct PredictorReg
    : public dmlc::FunctionRegEntryBase<PredictorReg, std::function<Predictor*(Context const*)>> {};
 
#define XGBOOST_REGISTER_PREDICTOR(UniqueId, Name)      \
  static DMLC_ATTRIBUTE_UNUSED ::xgboost::PredictorReg& \
      __make_##PredictorReg##_##UniqueId##__ =          \
          ::dmlc::Registry<::xgboost::PredictorReg>::Get()->__REGISTER__(Name)
}  // namespace xgboost