data_partition.hpp 8.26 KB
Newer Older
Guolin Ke's avatar
Guolin Ke committed
1
2
3
#ifndef LIGHTGBM_TREELEARNER_DATA_PARTITION_HPP_
#define LIGHTGBM_TREELEARNER_DATA_PARTITION_HPP_

Guolin Ke's avatar
Guolin Ke committed
4
#include <LightGBM/dataset.h>
5
#include <LightGBM/meta.h>
6
#include <LightGBM/utils/openmp_wrapper.h>
Guolin Ke's avatar
Guolin Ke committed
7

8
#include <algorithm>
Guolin Ke's avatar
Guolin Ke committed
9
10
11
12
13
14
15
16
#include <cstring>
#include <vector>

namespace LightGBM {
/*!
* \brief DataPartition is used to store the the partition of data on tree.
*/
class DataPartition {
Nikita Titov's avatar
Nikita Titov committed
17
 public:
18
19
  DataPartition(data_size_t num_data, int num_leaves)
    :num_data_(num_data), num_leaves_(num_leaves) {
Guolin Ke's avatar
Guolin Ke committed
20
21
22
23
24
    leaf_begin_.resize(num_leaves_);
    leaf_count_.resize(num_leaves_);
    indices_.resize(num_data_);
    temp_left_indices_.resize(num_data_);
    temp_right_indices_.resize(num_data_);
Guolin Ke's avatar
Guolin Ke committed
25
    used_data_indices_ = nullptr;
Guolin Ke's avatar
Guolin Ke committed
26
27
    #pragma omp parallel
    #pragma omp master
Guolin Ke's avatar
Guolin Ke committed
28
29
30
    {
      num_threads_ = omp_get_num_threads();
    }
Guolin Ke's avatar
Guolin Ke committed
31
32
33
34
35
    offsets_buf_.resize(num_threads_);
    left_cnts_buf_.resize(num_threads_);
    right_cnts_buf_.resize(num_threads_);
    left_write_pos_buf_.resize(num_threads_);
    right_write_pos_buf_.resize(num_threads_);
Guolin Ke's avatar
Guolin Ke committed
36
  }
37
38
39
40
41
42

  void ResetLeaves(int num_leaves) {
    num_leaves_ = num_leaves;
    leaf_begin_.resize(num_leaves_);
    leaf_count_.resize(num_leaves_);
  }
Guolin Ke's avatar
Guolin Ke committed
43
44
45
46
47
48
  void ResetNumData(int num_data) {
    num_data_ = num_data;
    indices_.resize(num_data_);
    temp_left_indices_.resize(num_data_);
    temp_right_indices_.resize(num_data_);
  }
Guolin Ke's avatar
Guolin Ke committed
49
50
51
52
53
54
55
  ~DataPartition() {
  }

  /*!
  * \brief Init, will put all data on the root(leaf_idx = 0)
  */
  void Init() {
56
57
    std::fill(leaf_begin_.begin(), leaf_begin_.end(), 0);
    std::fill(leaf_count_.begin(), leaf_count_.end(), 0);
Guolin Ke's avatar
Guolin Ke committed
58
59
60
    if (used_data_indices_ == nullptr) {
      // if using all data
      leaf_count_[0] = num_data_;
Guolin Ke's avatar
Guolin Ke committed
61
      #pragma omp parallel for schedule(static)
Guolin Ke's avatar
Guolin Ke committed
62
63
64
65
66
67
      for (data_size_t i = 0; i < num_data_; ++i) {
        indices_[i] = i;
      }
    } else {
      // if bagging
      leaf_count_[0] = used_data_count_;
Guolin Ke's avatar
Guolin Ke committed
68
      std::memcpy(indices_.data(), used_data_indices_, used_data_count_ * sizeof(data_size_t));
Guolin Ke's avatar
Guolin Ke committed
69
70
71
    }
  }

72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
  void ResetByLeafPred(const std::vector<int>& leaf_pred, int num_leaves) {
    ResetLeaves(num_leaves);
    std::vector<std::vector<data_size_t>> indices_per_leaf(num_leaves_);
    for (data_size_t i = 0; i < static_cast<data_size_t>(leaf_pred.size()); ++i) {
      indices_per_leaf[leaf_pred[i]].push_back(i);
    }
    data_size_t offset = 0;
    for (int i = 0; i < num_leaves_; ++i) {
      leaf_begin_[i] = offset;
      leaf_count_[i] = static_cast<data_size_t>(indices_per_leaf[i].size());
      std::copy(indices_per_leaf[i].begin(), indices_per_leaf[i].end(), indices_.begin() + leaf_begin_[i]);
      offset += leaf_count_[i];
    }
  }

Guolin Ke's avatar
Guolin Ke committed
87
88
89
90
91
92
  /*!
  * \brief Get the data indices of one leaf
  * \param leaf index of leaf
  * \param indices output data indices
  * \return number of data on this leaf
  */
Guolin Ke's avatar
Guolin Ke committed
93
  const data_size_t* GetIndexOnLeaf(int leaf, data_size_t* out_len) const {
Guolin Ke's avatar
Guolin Ke committed
94
95
    // copy reference, maybe unsafe, but faster
    data_size_t begin = leaf_begin_[leaf];
Guolin Ke's avatar
Guolin Ke committed
96
97
    *out_len = leaf_count_[leaf];
    return indices_.data() + begin;
Guolin Ke's avatar
Guolin Ke committed
98
99
100
101
102
103
104
105
106
  }

  /*!
  * \brief Split the data
  * \param leaf index of leaf
  * \param feature_bins feature bin data
  * \param threshold threshold that want to split
  * \param right_leaf index of right leaf
  */
107
  void Split(int leaf, const Dataset* dataset, int feature, const uint32_t* threshold, int num_threshold, bool default_left, int right_leaf) {
Guolin Ke's avatar
Guolin Ke committed
108
    const data_size_t min_inner_size = 512;
Guolin Ke's avatar
Guolin Ke committed
109
110
111
112
113
114
115
    // get leaf boundary
    const data_size_t begin = leaf_begin_[leaf];
    const data_size_t cnt = leaf_count_[leaf];

    data_size_t inner_size = (cnt + num_threads_ - 1) / num_threads_;
    if (inner_size < min_inner_size) { inner_size = min_inner_size; }
    // split data multi-threading
116
    OMP_INIT_EX();
Guolin Ke's avatar
Guolin Ke committed
117
    #pragma omp parallel for schedule(static, 1)
Guolin Ke's avatar
Guolin Ke committed
118
    for (int i = 0; i < num_threads_; ++i) {
119
      OMP_LOOP_EX_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
120
121
122
123
124
125
126
      left_cnts_buf_[i] = 0;
      right_cnts_buf_[i] = 0;
      data_size_t cur_start = i * inner_size;
      if (cur_start > cnt) { continue; }
      data_size_t cur_cnt = inner_size;
      if (cur_start + cur_cnt > cnt) { cur_cnt = cnt - cur_start; }
      // split data inner, reduce the times of function called
127
      data_size_t cur_left_count = dataset->Split(feature, threshold, num_threshold, default_left, indices_.data() + begin + cur_start, cur_cnt,
Guolin Ke's avatar
Guolin Ke committed
128
                                                  temp_left_indices_.data() + cur_start, temp_right_indices_.data() + cur_start);
Guolin Ke's avatar
Guolin Ke committed
129
130
131
      offsets_buf_[i] = cur_start;
      left_cnts_buf_[i] = cur_left_count;
      right_cnts_buf_[i] = cur_cnt - cur_left_count;
132
      OMP_LOOP_EX_END();
Guolin Ke's avatar
Guolin Ke committed
133
    }
134
    OMP_THROW_EX();
Guolin Ke's avatar
Guolin Ke committed
135
136
137
138
139
140
141
142
143
    data_size_t left_cnt = 0;
    left_write_pos_buf_[0] = 0;
    right_write_pos_buf_[0] = 0;
    for (int i = 1; i < num_threads_; ++i) {
      left_write_pos_buf_[i] = left_write_pos_buf_[i - 1] + left_cnts_buf_[i - 1];
      right_write_pos_buf_[i] = right_write_pos_buf_[i - 1] + right_cnts_buf_[i - 1];
    }
    left_cnt = left_write_pos_buf_[num_threads_ - 1] + left_cnts_buf_[num_threads_ - 1];
    // copy back indices of right leaf to indices_
Guolin Ke's avatar
Guolin Ke committed
144
    #pragma omp parallel for schedule(static, 1)
Guolin Ke's avatar
Guolin Ke committed
145
146
    for (int i = 0; i < num_threads_; ++i) {
      if (left_cnts_buf_[i] > 0) {
Guolin Ke's avatar
Guolin Ke committed
147
148
        std::memcpy(indices_.data() + begin + left_write_pos_buf_[i],
                    temp_left_indices_.data() + offsets_buf_[i], left_cnts_buf_[i] * sizeof(data_size_t));
Guolin Ke's avatar
Guolin Ke committed
149
150
      }
      if (right_cnts_buf_[i] > 0) {
Guolin Ke's avatar
Guolin Ke committed
151
152
        std::memcpy(indices_.data() + begin + left_cnt + right_write_pos_buf_[i],
                    temp_right_indices_.data() + offsets_buf_[i], right_cnts_buf_[i] * sizeof(data_size_t));
Guolin Ke's avatar
Guolin Ke committed
153
154
155
156
157
158
159
160
161
162
163
164
165
      }
    }
    // update leaf boundary
    leaf_count_[leaf] = left_cnt;
    leaf_begin_[right_leaf] = left_cnt + begin;
    leaf_count_[right_leaf] = cnt - left_cnt;
  }

  /*!
  * \brief SetLabelAt used data indices before training, used for bagging
  * \param used_data_indices indices of used data
  * \param num_used_data number of used data
  */
Guolin Ke's avatar
Guolin Ke committed
166
  void SetUsedDataIndices(const data_size_t* used_data_indices, data_size_t num_used_data) {
Guolin Ke's avatar
Guolin Ke committed
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
    used_data_indices_ = used_data_indices;
    used_data_count_ = num_used_data;
  }

  /*!
  * \brief Get number of data on one leaf
  * \param leaf index of leaf
  * \return number of data of this leaf
  */
  data_size_t leaf_count(int leaf) const { return leaf_count_[leaf]; }

  /*!
  * \brief Get leaf begin
  * \param leaf index of leaf
  * \return begin index of this leaf
  */
  data_size_t leaf_begin(int leaf) const { return leaf_begin_[leaf]; }

Guolin Ke's avatar
Guolin Ke committed
185
  const data_size_t* indices() const { return indices_.data(); }
Guolin Ke's avatar
Guolin Ke committed
186
187
188
189

  /*! \brief Get number of leaves */
  int num_leaves() const { return num_leaves_; }

Nikita Titov's avatar
Nikita Titov committed
190
 private:
Guolin Ke's avatar
Guolin Ke committed
191
192
193
194
195
  /*! \brief Number of all data */
  data_size_t num_data_;
  /*! \brief Number of all leaves */
  int num_leaves_;
  /*! \brief start index of data on one leaf */
Guolin Ke's avatar
Guolin Ke committed
196
  std::vector<data_size_t> leaf_begin_;
Guolin Ke's avatar
Guolin Ke committed
197
  /*! \brief number of data on one leaf */
Guolin Ke's avatar
Guolin Ke committed
198
  std::vector<data_size_t> leaf_count_;
Guolin Ke's avatar
Guolin Ke committed
199
  /*! \brief Store all data's indices, order by leaf[data_in_leaf0,..,data_leaf1,..] */
Guolin Ke's avatar
Guolin Ke committed
200
  std::vector<data_size_t> indices_;
Guolin Ke's avatar
Guolin Ke committed
201
  /*! \brief team indices buffer for split */
Guolin Ke's avatar
Guolin Ke committed
202
  std::vector<data_size_t> temp_left_indices_;
Guolin Ke's avatar
Guolin Ke committed
203
  /*! \brief team indices buffer for split */
Guolin Ke's avatar
Guolin Ke committed
204
  std::vector<data_size_t> temp_right_indices_;
Guolin Ke's avatar
Guolin Ke committed
205
206
207
208
209
210
211
  /*! \brief used data indices, used for bagging */
  const data_size_t* used_data_indices_;
  /*! \brief used data count, used for bagging */
  data_size_t used_data_count_;
  /*! \brief number of threads */
  int num_threads_;
  /*! \brief Buffer for multi-threading data partition, used to store offset for different threads */
Guolin Ke's avatar
Guolin Ke committed
212
  std::vector<data_size_t> offsets_buf_;
Guolin Ke's avatar
Guolin Ke committed
213
  /*! \brief Buffer for multi-threading data partition, used to store left count after split for different threads */
Guolin Ke's avatar
Guolin Ke committed
214
  std::vector<data_size_t> left_cnts_buf_;
Guolin Ke's avatar
Guolin Ke committed
215
  /*! \brief Buffer for multi-threading data partition, used to store right count after split for different threads */
Guolin Ke's avatar
Guolin Ke committed
216
  std::vector<data_size_t> right_cnts_buf_;
Guolin Ke's avatar
Guolin Ke committed
217
  /*! \brief Buffer for multi-threading data partition, used to store write position of left leaf for different threads */
Guolin Ke's avatar
Guolin Ke committed
218
  std::vector<data_size_t> left_write_pos_buf_;
Guolin Ke's avatar
Guolin Ke committed
219
  /*! \brief Buffer for multi-threading data partition, used to store write position of right leaf for different threads */
Guolin Ke's avatar
Guolin Ke committed
220
  std::vector<data_size_t> right_write_pos_buf_;
Guolin Ke's avatar
Guolin Ke committed
221
222
223
};

}  // namespace LightGBM
Guolin Ke's avatar
Guolin Ke committed
224
#endif   // LightGBM_TREELEARNER_DATA_PARTITION_HPP_