tree_updater.h

Go to the documentation of this file.

 
#ifndef XGBOOST_TREE_UPDATER_H_
#define XGBOOST_TREE_UPDATER_H_
 
#include <dmlc/registry.h>
#include <xgboost/base.h>                // for Args, GradientPair
#include <xgboost/data.h>                // for DMatrix
#include <xgboost/gradient.h>            // for GradientContainer
#include <xgboost/host_device_vector.h>  // for HostDeviceVector
#include <xgboost/linalg.h>              // for VectorView
#include <xgboost/model.h>               // for Configurable
#include <xgboost/span.h>                // for Span
#include <xgboost/tree_model.h>          // for RegTree
 
#include <functional>  // for function
#include <string>      // for string
#include <vector>      // for 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; }
  virtual void Update(tree::TrainParam const* param, GradientContainer* gpair, DMatrix* p_fmat,
                      common::Span<HostDeviceVector<bst_node_t>> out_position,
                      std::vector<RegTree*> const& out_trees) = 0;
 
  [[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_