multi_target_tree_model.h

前往此文件文档。

 
#ifndef XGBOOST_MULTI_TARGET_TREE_MODEL_H_
#define XGBOOST_MULTI_TARGET_TREE_MODEL_H_
 
#include <xgboost/base.h> // for bst_node_t, bst_target_t, bst_feature_t
#include <xgboost/context.h> // for Context
#include <xgboost/host_device_vector.h> // for HostDeviceVector
#include <xgboost/linalg.h> // for VectorView, MatrixView
#include <xgboost/model.h> // for Model
#include <xgboost/span.h> // for Span
 
#include <cstddef> // for size_t
#include <cstdint> // for uint8_t
#include <mutex> // for mutex
#include <vector> // for vector
 
namespace xgboost {
struct TreeParam;
struct MultiTargetTreeView {
 static bst_node_t constexpr InvalidNodeId() { return -1; }
 
 bst_node_t const* left;
 bst_node_t const* right;
 bst_node_t const* parent;
 
 bst_feature_t const* split_index;
 std::uint8_t const* default_left;
 float const* split_conds;
 
 // The number of nodes
 std::size_t n{0};
 
 linalg::MatrixView<float const> weights;
 
 [[nodiscard]] XGBOOST_DEVICE bool IsLeaf(bst_node_t nidx) const {
 return left[nidx] == InvalidNodeId();
  }
 
 [[nodiscard]] XGBOOST_DEVICE bst_node_t LeftChild(bst_node_t nidx) const { return left[nidx]; }
 [[nodiscard]] XGBOOST_DEVICE bst_node_t RightChild(bst_node_t nidx) const { return right[nidx]; }
 [[nodiscard]] XGBOOST_DEVICE bst_feature_t SplitIndex(bst_node_t nidx) const {
 return split_index[nidx];
  }
 [[nodiscard]] XGBOOST_DEVICE float SplitCond(bst_node_t nidx) const { return split_conds[nidx]; }
 [[nodiscard]] XGBOOST_DEVICE bool DefaultLeft(bst_node_t nidx) const {
 return default_left[nidx];
  }
 [[nodiscard]] XGBOOST_DEVICE bst_node_t DefaultChild(bst_node_t nidx) const {
 return this->DefaultLeft(nidx) ? this->LeftChild(nidx) : this->RightChild(nidx);
  }
 [[nodiscard]] XGBOOST_DEVICE linalg::VectorView<float const> LeafValue(bst_node_t nidx) const {
 return this->weights.Slice(nidx, linalg::All());
  }
 
 [[nodiscard]] bst_target_t NumTargets() const { return this->weights.Shape(1); }
 [[nodiscard]] bst_node_t Size() const { return this->n; }
};
 
class MultiTargetTree : public Model {
 public
 static bst_node_t constexpr InvalidNodeId() { return MultiTargetTreeView::InvalidNodeId(); }
 
 private
 TreeParam const* param_;
 HostDeviceVector<bst_node_t> left_;
 HostDeviceVector<bst_node_t> right_;
 HostDeviceVector<bst_node_t> parent_;
 HostDeviceVector<bst_feature_t> split_index_;
 HostDeviceVector<std::uint8_t> default_left_;
 HostDeviceVector<float> split_conds_;
 HostDeviceVector<float> weights_;
 
 mutable std::mutex tree_view_lock_;
 
 [[nodiscard]] linalg::VectorView<float const> NodeWeight(bst_node_t nidx) const {
 auto beg = nidx * this->NumTargets();
 auto v = this->weights_.ConstHostSpan().subspan(beg, this->NumTargets());
 return linalg::MakeTensorView(DeviceOrd::CPU(), v, v.size());
  }
 [[nodiscard]] linalg::VectorView<float> NodeWeight(bst_node_t nidx) {
 auto beg = nidx * this->NumTargets();
 auto v = this->weights_.HostSpan().subspan(beg, this->NumTargets());
 return linalg::MakeTensorView(DeviceOrd::CPU(), v, v.size());
  }
 
 public
 explicit MultiTargetTree(TreeParam const* param);
 MultiTargetTree(MultiTargetTree const& that);
 MultiTargetTree& operator=(MultiTargetTree const& that) = delete;
 MultiTargetTree(MultiTargetTree&& that) = delete;
 MultiTargetTree& operator=(MultiTargetTree&& that) = delete;
 
 void SetLeaf(bst_node_t nidx, linalg::VectorView<float const> weight);
 void Expand(bst_node_t nidx, bst_feature_t split_idx, float split_cond, bool default_left,
 linalg::VectorView<float const> base_weight,
 linalg::VectorView<float const> left_weight,
 linalg::VectorView<float const> right_weight);
 
 [[nodiscard]] bool IsLeaf(bst_node_t nidx) const {
 return left_.ConstHostVector()[nidx] == InvalidNodeId();
  }
 [[nodiscard]] bst_node_t Parent(bst_node_t nidx) const {
 return parent_.ConstHostVector().at(nidx);
  }
 [[nodiscard]] bst_node_t LeftChild(bst_node_t nidx) const {
 return left_.ConstHostVector().at(nidx);
  }
 [[nodiscard]] bst_node_t RightChild(bst_node_t nidx) const {
 return right_.ConstHostVector().at(nidx);
  }
 
 [[nodiscard]] bst_feature_t SplitIndex(bst_node_t nidx) const {
 return split_index_.ConstHostVector()[nidx];
  }
 [[nodiscard]] float SplitCond(bst_node_t nidx) const {
 return split_conds_.ConstHostVector()[nidx];
  }
 [[nodiscard]] bool DefaultLeft(bst_node_t nidx) const {
 return default_left_.ConstHostVector()[nidx];
  }
 [[nodiscard]] bst_node_t DefaultChild(bst_node_t nidx) const {
 return this->DefaultLeft(nidx) ? this->LeftChild(nidx) : this->RightChild(nidx);
  }
 
 [[nodiscard]] bst_target_t NumTargets() const;
 
 [[nodiscard]] std::size_t Size() const;
 
 [[nodiscard]] bst_node_t Depth(bst_node_t nidx) const {
 bst_node_t depth{0};
 while (Parent(nidx) != InvalidNodeId()) {
 ++depth;
 nidx = Parent(nidx);
    }
 return depth;
  }
 
 [[nodiscard]] linalg::VectorView<float const> LeafValue(bst_node_t nidx) const {
 CHECK(IsLeaf(nidx));
 return this->NodeWeight(nidx);
  }
 [[nodiscard]] MultiTargetTreeView View(Context const* ctx) const;
 
 void LoadModel(Json const& in) override;
 void SaveModel(Json* out) const override;
 
 [[nodiscard]] std::size_t MemCostBytes() const;
};
} // namespace xgboost
#endif // XGBOOST_MULTI_TARGET_TREE_MODEL_H_