predictor.h

前往此文件的文档。

#pragma once
#include <xgboost/base.h>
#include <xgboost/cache.h> // 用于 DMatrixCache
#include <xgboost/context.h> // 用于 Context
#include <xgboost/context.h>
#include <xgboost/data.h>
#include <xgboost/host_device_vector.h>

#include <functional> // 用于 function
#include <memory> // 用于 shared_ptr
#include <string>
#include <vector>

// 前向声明
namespace xgboost::gbm {
struct GBTreeModel;
} // namespace xgboost::gbm

namespace xgboost {
struct PredictionCacheEntry {
 // 用于缓存预测值的存储
 HostDeviceVector<float> predictions;
 // 当前缓存的版本，对应于树的层数
 std::uint32_t version{0};

 PredictionCacheEntry() = default;
 void Update(std::uint32_t v) { version += v; }
 void Reset() { version = 0; }
};

class PredictionContainer : public DMatrixCache<PredictionCacheEntry> {
 // 我们为所有线程缓存最多 64 个 DMatrix
 std::size_t static constexpr DefaultSize() { return 64; }

 public
 PredictionContainer() : DMatrixCache<PredictionCacheEntry>{DefaultSize()} {}
 std::shared_ptr<PredictionCacheEntry> Cache(std::shared_ptr<DMatrix> m, DeviceOrd device) {
 auto p_cache = this->CacheItem(m);
 if (!device.IsCPU()) {
 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&);

 virtual void InitOutPredictions(const MetaInfo& info, HostDeviceVector<float>* out_predt,
 const gbm::GBTreeModel& model) const;

 virtual void PredictBatch(DMatrix* dmat, PredictionCacheEntry* out_preds,
 gbm::GBTreeModel const& model, bst_tree_t tree_begin,
 bst_tree_t tree_end = 0) const = 0;

 virtual bool InplacePredict(std::shared_ptr<DMatrix> p_fmat, const gbm::GBTreeModel& model,
 float missing, PredictionCacheEntry* out_preds,
 bst_tree_t tree_begin = 0, bst_tree_t tree_end = 0) const = 0;

 virtual void PredictLeaf(DMatrix* dmat, HostDeviceVector<float>* out_preds,
 gbm::GBTreeModel const& model, bst_tree_t tree_end = 0) const = 0;

 virtual void PredictContribution(DMatrix* dmat, HostDeviceVector<float>* out_contribs,
 gbm::GBTreeModel const& model, bst_tree_t tree_end = 0,
 std::vector<float> const* tree_weights = nullptr,
 bool approximate = false, int condition = 0,
 unsigned condition_feature = 0) const = 0;

 virtual void PredictInteractionContributions(DMatrix* dmat, HostDeviceVector<float>* out_contribs,
 gbm::GBTreeModel const& model,
 bst_tree_t tree_end = 0,
 std::vector<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