tree_model.h

Go to the documentation of this file.

 
#ifndef XGBOOST_TREE_MODEL_H_
#define XGBOOST_TREE_MODEL_H_
 
#include <dmlc/io.h>
#include <dmlc/parameter.h>
 
#include <xgboost/base.h>
#include <xgboost/data.h>
#include <xgboost/logging.h>
#include <xgboost/feature_map.h>
#include <xgboost/model.h>
 
#include <limits>
#include <vector>
#include <string>
#include <cstring>
#include <algorithm>
#include <tuple>
#include <stack>
 
namespace xgboost {
 
struct PathElement;  // forward declaration
 
class Json;
// FIXME(trivialfis): Once binary IO is gone, make this parameter internal as it should
// not be configured by users.
struct TreeParam : public dmlc::Parameter<TreeParam> {
  int deprecated_num_roots;
  int num_nodes;
  int num_deleted;
  int deprecated_max_depth;
  bst_feature_t num_feature;
  int size_leaf_vector;
  int reserved[31];
  TreeParam() {
    // assert compact alignment
    static_assert(sizeof(TreeParam) == (31 + 6) * sizeof(int),
                  "TreeParam: 64 bit align");
    std::memset(this, 0, sizeof(TreeParam));
    num_nodes = 1;
    deprecated_num_roots = 1;
  }
 
  // Swap byte order for all fields. Useful for transporting models between machines with different
  // endianness (big endian vs little endian)
  inline TreeParam ByteSwap() const {
    TreeParam x = *this;
    dmlc::ByteSwap(&x.deprecated_num_roots, sizeof(x.deprecated_num_roots), 1);
    dmlc::ByteSwap(&x.num_nodes, sizeof(x.num_nodes), 1);
    dmlc::ByteSwap(&x.num_deleted, sizeof(x.num_deleted), 1);
    dmlc::ByteSwap(&x.deprecated_max_depth, sizeof(x.deprecated_max_depth), 1);
    dmlc::ByteSwap(&x.num_feature, sizeof(x.num_feature), 1);
    dmlc::ByteSwap(&x.size_leaf_vector, sizeof(x.size_leaf_vector), 1);
    dmlc::ByteSwap(x.reserved, sizeof(x.reserved[0]), sizeof(x.reserved) / sizeof(x.reserved[0]));
    return x;
  }
 
  // declare the parameters
  DMLC_DECLARE_PARAMETER(TreeParam) {
    // only declare the parameters that can be set by the user.
    // other arguments are set by the algorithm.
    DMLC_DECLARE_FIELD(num_nodes).set_lower_bound(1).set_default(1);
    DMLC_DECLARE_FIELD(num_feature)
        .describe("Number of features used in tree construction.");
    DMLC_DECLARE_FIELD(num_deleted);
    DMLC_DECLARE_FIELD(size_leaf_vector).set_lower_bound(0).set_default(0)
        .describe("Size of leaf vector, reserved for vector tree");
  }
 
  bool operator==(const TreeParam& b) const {
    return num_nodes == b.num_nodes &&
           num_deleted == b.num_deleted &&
           num_feature == b.num_feature &&
           size_leaf_vector == b.size_leaf_vector;
  }
};
 
struct RTreeNodeStat {
  bst_float loss_chg;
  bst_float sum_hess;
  bst_float base_weight;
  int leaf_child_cnt {0};
 
  RTreeNodeStat() = default;
  RTreeNodeStat(float loss_chg, float sum_hess, float weight) :
      loss_chg{loss_chg}, sum_hess{sum_hess}, base_weight{weight} {}
  bool operator==(const RTreeNodeStat& b) const {
    return loss_chg == b.loss_chg && sum_hess == b.sum_hess &&
           base_weight == b.base_weight && leaf_child_cnt == b.leaf_child_cnt;
  }
  // Swap byte order for all fields. Useful for transporting models between machines with different
  // endianness (big endian vs little endian)
  inline RTreeNodeStat ByteSwap() const {
    RTreeNodeStat x = *this;
    dmlc::ByteSwap(&x.loss_chg, sizeof(x.loss_chg), 1);
    dmlc::ByteSwap(&x.sum_hess, sizeof(x.sum_hess), 1);
    dmlc::ByteSwap(&x.base_weight, sizeof(x.base_weight), 1);
    dmlc::ByteSwap(&x.leaf_child_cnt, sizeof(x.leaf_child_cnt), 1);
    return x;
  }
};
 
class RegTree : public Model {
 public:
  using SplitCondT = bst_float;
  static constexpr bst_node_t kInvalidNodeId {-1};
  static constexpr uint32_t kDeletedNodeMarker = std::numeric_limits<uint32_t>::max();
  static constexpr bst_node_t kRoot { 0 };
 
  class Node {
   public:
    XGBOOST_DEVICE Node()  {
      // assert compact alignment
      static_assert(sizeof(Node) == 4 * sizeof(int) + sizeof(Info),
                    "Node: 64 bit align");
    }
    Node(int32_t cleft, int32_t cright, int32_t parent,
         uint32_t split_ind, float split_cond, bool default_left) :
        parent_{parent}, cleft_{cleft}, cright_{cright} {
      this->SetParent(parent_);
      this->SetSplit(split_ind, split_cond, default_left);
    }
 
    XGBOOST_DEVICE int LeftChild() const {
      return this->cleft_;
    }
    XGBOOST_DEVICE int RightChild() const {
      return this->cright_;
    }
    XGBOOST_DEVICE int DefaultChild() const {
      return this->DefaultLeft() ? this->LeftChild() : this->RightChild();
    }
    XGBOOST_DEVICE unsigned SplitIndex() const {
      return sindex_ & ((1U << 31) - 1U);
    }
    XGBOOST_DEVICE bool DefaultLeft() const {
      return (sindex_ >> 31) != 0;
    }
    XGBOOST_DEVICE bool IsLeaf() const {
      return cleft_ == kInvalidNodeId;
    }
    XGBOOST_DEVICE bst_float LeafValue() const {
      return (this->info_).leaf_value;
    }
    XGBOOST_DEVICE SplitCondT SplitCond() const {
      return (this->info_).split_cond;
    }
    XGBOOST_DEVICE int Parent() const {
      return parent_ & ((1U << 31) - 1);
    }
    XGBOOST_DEVICE bool IsLeftChild() const {
      return (parent_ & (1U << 31)) != 0;
    }
    XGBOOST_DEVICE bool IsDeleted() const {
      return sindex_ == kDeletedNodeMarker;
    }
    XGBOOST_DEVICE bool IsRoot() const { return parent_ == kInvalidNodeId; }
    XGBOOST_DEVICE void SetLeftChild(int nid) {
      this->cleft_ = nid;
    }
    XGBOOST_DEVICE void SetRightChild(int nid) {
      this->cright_ = nid;
    }
    XGBOOST_DEVICE void SetSplit(unsigned split_index, SplitCondT split_cond,
                          bool default_left = false) {
      if (default_left) split_index |= (1U << 31);
      this->sindex_ = split_index;
      (this->info_).split_cond = split_cond;
    }
    XGBOOST_DEVICE void SetLeaf(bst_float value, int right = kInvalidNodeId) {
      (this->info_).leaf_value = value;
      this->cleft_ = kInvalidNodeId;
      this->cright_ = right;
    }
    XGBOOST_DEVICE void MarkDelete() {
      this->sindex_ = kDeletedNodeMarker;
    }
    XGBOOST_DEVICE void Reuse() {
      this->sindex_ = 0;
    }
    // set parent
    XGBOOST_DEVICE void SetParent(int pidx, bool is_left_child = true) {
      if (is_left_child) pidx |= (1U << 31);
      this->parent_ = pidx;
    }
    bool operator==(const Node& b) const {
      return parent_ == b.parent_ && cleft_ == b.cleft_ &&
             cright_ == b.cright_ && sindex_ == b.sindex_ &&
             info_.leaf_value == b.info_.leaf_value;
    }
 
    inline Node ByteSwap() const {
      Node x = *this;
      dmlc::ByteSwap(&x.parent_, sizeof(x.parent_), 1);
      dmlc::ByteSwap(&x.cleft_, sizeof(x.cleft_), 1);
      dmlc::ByteSwap(&x.cright_, sizeof(x.cright_), 1);
      dmlc::ByteSwap(&x.sindex_, sizeof(x.sindex_), 1);
      dmlc::ByteSwap(&x.info_, sizeof(x.info_), 1);
      return x;
    }
 
   private:
    union Info{
      bst_float leaf_value;
      SplitCondT split_cond;
    };
    // pointer to parent, highest bit is used to
    // indicate whether it's a left child or not
    int32_t parent_{kInvalidNodeId};
    // pointer to left, right
    int32_t cleft_{kInvalidNodeId}, cright_{kInvalidNodeId};
    // split feature index, left split or right split depends on the highest bit
    uint32_t sindex_{0};
    // extra info
    Info info_;
  };
 
  void ChangeToLeaf(int rid, bst_float value) {
    CHECK(nodes_[nodes_[rid].LeftChild() ].IsLeaf());
    CHECK(nodes_[nodes_[rid].RightChild()].IsLeaf());
    this->DeleteNode(nodes_[rid].LeftChild());
    this->DeleteNode(nodes_[rid].RightChild());
    nodes_[rid].SetLeaf(value);
  }
  void CollapseToLeaf(int rid, bst_float value) {
    if (nodes_[rid].IsLeaf()) return;
    if (!nodes_[nodes_[rid].LeftChild() ].IsLeaf()) {
      CollapseToLeaf(nodes_[rid].LeftChild(), 0.0f);
    }
    if (!nodes_[nodes_[rid].RightChild() ].IsLeaf()) {
      CollapseToLeaf(nodes_[rid].RightChild(), 0.0f);
    }
    this->ChangeToLeaf(rid, value);
  }
 
  TreeParam param;
  RegTree() {
    param.num_nodes = 1;
    param.num_deleted = 0;
    nodes_.resize(param.num_nodes);
    stats_.resize(param.num_nodes);
    split_types_.resize(param.num_nodes, FeatureType::kNumerical);
    split_categories_segments_.resize(param.num_nodes);
    for (int i = 0; i < param.num_nodes; i ++) {
      nodes_[i].SetLeaf(0.0f);
      nodes_[i].SetParent(kInvalidNodeId);
    }
  }
  Node& operator[](int nid) {
    return nodes_[nid];
  }
  const Node& operator[](int nid) const {
    return nodes_[nid];
  }
 
  const std::vector<Node>& GetNodes() const { return nodes_; }
 
  const std::vector<RTreeNodeStat>& GetStats() const { return stats_; }
 
  RTreeNodeStat& Stat(int nid) {
    return stats_[nid];
  }
  const RTreeNodeStat& Stat(int nid) const {
    return stats_[nid];
  }
 
  void Load(dmlc::Stream* fi);
  void Save(dmlc::Stream* fo) const;
 
  void LoadModel(Json const& in) override;
  void SaveModel(Json* out) const override;
 
  bool operator==(const RegTree& b) const {
    return nodes_ == b.nodes_ && stats_ == b.stats_ &&
           deleted_nodes_ == b.deleted_nodes_ && param == b.param;
  }
  /* \brief Iterate through all nodes in this tree.
   *
   * \param Function that accepts a node index, and returns false when iteration should
   *        stop, otherwise returns true.
   */
  template <typename Func> void WalkTree(Func func) const {
    std::stack<bst_node_t> nodes;
    nodes.push(kRoot);
    auto &self = *this;
    while (!nodes.empty()) {
      auto nidx = nodes.top();
      nodes.pop();
      if (!func(nidx)) {
        return;
      }
      auto left = self[nidx].LeftChild();
      auto right = self[nidx].RightChild();
      if (left != RegTree::kInvalidNodeId) {
        nodes.push(left);
      }
      if (right != RegTree::kInvalidNodeId) {
        nodes.push(right);
      }
    }
  }
  bool Equal(const RegTree& b) const;
 
  void ExpandNode(bst_node_t nid, unsigned split_index, bst_float split_value,
                  bool default_left, bst_float base_weight,
                  bst_float left_leaf_weight, bst_float right_leaf_weight,
                  bst_float loss_change, float sum_hess, float left_sum,
                  float right_sum,
                  bst_node_t leaf_right_child = kInvalidNodeId);
 
  void ExpandCategorical(bst_node_t nid, unsigned split_index,
                         common::Span<const uint32_t> split_cat, bool default_left,
                         bst_float base_weight, bst_float left_leaf_weight,
                         bst_float right_leaf_weight, bst_float loss_change, float sum_hess,
                         float left_sum, float right_sum);
 
  bool HasCategoricalSplit() const {
    return !split_categories_.empty();
  }
 
  int GetDepth(int nid) const {
    int depth = 0;
    while (!nodes_[nid].IsRoot()) {
      ++depth;
      nid = nodes_[nid].Parent();
    }
    return depth;
  }
 
  int MaxDepth(int nid) const {
    if (nodes_[nid].IsLeaf()) return 0;
    return std::max(MaxDepth(nodes_[nid].LeftChild())+1,
                     MaxDepth(nodes_[nid].RightChild())+1);
  }
 
  int MaxDepth() {
    return MaxDepth(0);
  }
 
  int NumExtraNodes() const {
    return param.num_nodes - 1 - param.num_deleted;
  }
 
  /* \brief Count number of leaves in tree. */
  bst_node_t GetNumLeaves() const;
  bst_node_t GetNumSplitNodes() const;
 
  struct FVec {
    void Init(size_t size);
    void Fill(const SparsePage::Inst& inst);
 
    void Drop();
    size_t Size() const;
    bst_float GetFvalue(size_t i) const;
    bool IsMissing(size_t i) const;
    bool HasMissing() const;
 
 
   private:
    union Entry {
      bst_float fvalue;
      int flag;
    };
    std::vector<Entry> data_;
    bool has_missing_;
  };
 
  void CalculateContributions(const RegTree::FVec& feat,
                              std::vector<float>* mean_values,
                              bst_float* out_contribs, int condition = 0,
                              unsigned condition_feature = 0) const;
  void TreeShap(const RegTree::FVec& feat, bst_float* phi, bst_node_t node_index,
                unsigned unique_depth, PathElement* parent_unique_path,
                bst_float parent_zero_fraction, bst_float parent_one_fraction,
                int parent_feature_index, int condition,
                unsigned condition_feature, bst_float condition_fraction) const;
 
  void CalculateContributionsApprox(const RegTree::FVec& feat,
                                    std::vector<float>* mean_values,
                                    bst_float* out_contribs) const;
  std::string DumpModel(const FeatureMap& fmap,
                        bool with_stats,
                        std::string format) const;
  FeatureType NodeSplitType(bst_node_t nidx) const {
    return split_types_.at(nidx);
  }
  std::vector<FeatureType> const &GetSplitTypes() const { return split_types_; }
  common::Span<uint32_t const> GetSplitCategories() const { return split_categories_; }
  common::Span<uint32_t const> NodeCats(bst_node_t nidx) const {
    auto node_ptr = GetCategoriesMatrix().node_ptr;
    auto categories = GetCategoriesMatrix().categories;
    auto segment = node_ptr[nidx];
    auto node_cats = categories.subspan(segment.beg, segment.size);
    return node_cats;
  }
  auto const& GetSplitCategoriesPtr() const { return split_categories_segments_; }
 
  // The fields of split_categories_segments_[i] are set such that
  // the range split_categories_[beg:(beg+size)] stores the bitset for
  // the matching categories for the i-th node.
  struct Segment {
    size_t beg {0};
    size_t size {0};
  };
 
  struct CategoricalSplitMatrix {
    common::Span<FeatureType const> split_type;
    common::Span<uint32_t const> categories;
    common::Span<Segment const> node_ptr;
  };
 
  CategoricalSplitMatrix GetCategoriesMatrix() const {
    CategoricalSplitMatrix view;
    view.split_type = common::Span<FeatureType const>(this->GetSplitTypes());
    view.categories = this->GetSplitCategories();
    view.node_ptr = common::Span<Segment const>(split_categories_segments_);
    return view;
  }
 
 private:
  template <bool typed>
  void LoadCategoricalSplit(Json const& in);
  void SaveCategoricalSplit(Json* p_out) const;
  // vector of nodes
  std::vector<Node> nodes_;
  // free node space, used during training process
  std::vector<int>  deleted_nodes_;
  // stats of nodes
  std::vector<RTreeNodeStat> stats_;
  std::vector<FeatureType> split_types_;
 
  // Categories for each internal node.
  std::vector<uint32_t> split_categories_;
  // Ptr to split categories of each node.
  std::vector<Segment> split_categories_segments_;
 
  // allocate a new node,
  // !!!!!! NOTE: may cause BUG here, nodes.resize
  bst_node_t AllocNode() {
    if (param.num_deleted != 0) {
      int nid = deleted_nodes_.back();
      deleted_nodes_.pop_back();
      nodes_[nid].Reuse();
      --param.num_deleted;
      return nid;
    }
    int nd = param.num_nodes++;
    CHECK_LT(param.num_nodes, std::numeric_limits<int>::max())
        << "number of nodes in the tree exceed 2^31";
    nodes_.resize(param.num_nodes);
    stats_.resize(param.num_nodes);
    split_types_.resize(param.num_nodes, FeatureType::kNumerical);
    split_categories_segments_.resize(param.num_nodes);
    return nd;
  }
  // delete a tree node, keep the parent field to allow trace back
  void DeleteNode(int nid) {
    CHECK_GE(nid, 1);
    auto pid = (*this)[nid].Parent();
    if (nid == (*this)[pid].LeftChild()) {
      (*this)[pid].SetLeftChild(kInvalidNodeId);
    } else {
      (*this)[pid].SetRightChild(kInvalidNodeId);
    }
 
    deleted_nodes_.push_back(nid);
    nodes_[nid].MarkDelete();
    ++param.num_deleted;
  }
};
 
inline void RegTree::FVec::Init(size_t size) {
  Entry e; e.flag = -1;
  data_.resize(size);
  std::fill(data_.begin(), data_.end(), e);
  has_missing_ = true;
}
 
inline void RegTree::FVec::Fill(const SparsePage::Inst& inst) {
  size_t feature_count = 0;
  for (auto const& entry : inst) {
    if (entry.index >= data_.size()) {
      continue;
    }
    data_[entry.index].fvalue = entry.fvalue;
    ++feature_count;
  }
  has_missing_ = data_.size() != feature_count;
}
 
inline void RegTree::FVec::Drop() {
  Entry e{};
  e.flag = -1;
  std::fill_n(data_.data(), data_.size(), e);
  has_missing_ = true;
}
 
inline size_t RegTree::FVec::Size() const {
  return data_.size();
}
 
inline bst_float RegTree::FVec::GetFvalue(size_t i) const {
  return data_[i].fvalue;
}
 
inline bool RegTree::FVec::IsMissing(size_t i) const {
  return data_[i].flag == -1;
}
 
inline bool RegTree::FVec::HasMissing() const {
  return has_missing_;
}
}  // namespace xgboost
#endif  // XGBOOST_TREE_MODEL_H_