tree_updater.h

前往此文件文档。

 
#ifndef XGBOOST_TREE_UPDATER_H_
#define XGBOOST_TREE_UPDATER_H_
 
#include <dmlc/registry.h>
#include <xgboost/base.h> // 用于 Args, GradientPair
#include <xgboost/data.h> // DMatrix
#include <xgboost/host_device_vector.h> // 用于 HostDeviceVector
#include <xgboost/linalg.h> // 用于 VectorView
#include <xgboost/model.h> // 用于 Configurable
#include <xgboost/span.h> // 用于 Span
#include <xgboost/tree_model.h> // 用于 RegTree
 
#include <functional> // 用于 function
#include <string> // 用于 string
#include <vector> // 用于 vector
 
namespace xgboost {
namespace tree {
struct TrainParam;
}
 
class Json;
struct Context;
struct ObjInfo;
 
class TreeUpdater : public Configurable {
 protected
 Context const* ctx_ = nullptr;
 
 public
 explicit TreeUpdater(const Context* ctx) : ctx_(ctx) {}
 ~TreeUpdater() override = default;
 virtual void Configure(const Args& args) = 0;
 [[nodiscard]] virtual bool CanModifyTree() const { return false; }
 [[nodiscard]] virtual bool HasNodePosition() const { return false; }
 virtual void Update(tree::TrainParam const* param, linalg::Matrix<GradientPair>* gpair,
 DMatrix* data, common::Span<HostDeviceVector<bst_node_t>> out_position,
 const std::vector<RegTree*>& out_trees) = 0;
 
 virtual bool UpdatePredictionCache(const DMatrix* /*data*/,
 linalg::MatrixView<float> /*out_preds*/) {
 return false;
  }
 
 [[nodiscard]] virtual char const* Name() const = 0;
 
 static TreeUpdater* Create(const std::string& name, Context const* ctx, ObjInfo const* task);
};
 
struct TreeUpdaterReg
 : public dmlc::FunctionRegEntryBase<
 TreeUpdaterReg, std::function<TreeUpdater*(Context const* ctx, ObjInfo const* task)>> {};
 
#define XGBOOST_REGISTER_TREE_UPDATER(UniqueId, Name) \
 static DMLC_ATTRIBUTE_UNUSED ::xgboost::TreeUpdaterReg& \
 __make_ ## TreeUpdaterReg ## _ ## UniqueId ## __ = \
 ::dmlc::Registry< ::xgboost::TreeUpdaterReg>::Get()->__REGISTER__(Name)
 
} // namespace xgboost
#endif // XGBOOST_TREE_UPDATER_H_