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>
#include <xgboost/data.h>
#include <xgboost/generic_parameters.h>
#include <xgboost/host_device_vector.h>
#include <xgboost/linalg.h>
#include <xgboost/model.h>
#include <xgboost/task.h>
#include <xgboost/tree_model.h>
 
#include <functional>
#include <string>
#include <utility>
#include <vector>
 
namespace xgboost {
 
class Json;
 
class TreeUpdater : public Configurable {
 protected:
  GenericParameter const* ctx_ = nullptr;
 
 public:
  explicit TreeUpdater(const GenericParameter* ctx) : ctx_(ctx) {}
  ~TreeUpdater() override = default;
  virtual void Configure(const Args& args) = 0;
  virtual bool CanModifyTree() const { return false; }
  virtual bool HasNodePosition() const { return false; }
  virtual void Update(HostDeviceVector<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::VectorView<float> /*out_preds*/) {
    return false;
  }
 
  virtual char const* Name() const = 0;
 
  static TreeUpdater* Create(const std::string& name, GenericParameter const* tparam, ObjInfo task);
};
 
struct TreeUpdaterReg
    : public dmlc::FunctionRegEntryBase<
          TreeUpdaterReg,
          std::function<TreeUpdater*(GenericParameter const* tparam, ObjInfo 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_