data_partition.hpp 8.78 KB
Newer Older
1
2
3
4
/*!
 * Copyright (c) 2016 Microsoft Corporation. All rights reserved.
 * Licensed under the MIT License. See LICENSE file in the project root for license information.
 */
Guolin Ke's avatar
Guolin Ke committed
5
6
7
#ifndef LIGHTGBM_TREELEARNER_DATA_PARTITION_HPP_
#define LIGHTGBM_TREELEARNER_DATA_PARTITION_HPP_

Guolin Ke's avatar
Guolin Ke committed
8
#include <LightGBM/dataset.h>
9
#include <LightGBM/meta.h>
10
#include <LightGBM/utils/openmp_wrapper.h>
Guolin Ke's avatar
Guolin Ke committed
11
#include <LightGBM/utils/threading.h>
Guolin Ke's avatar
Guolin Ke committed
12

13
#include <algorithm>
Guolin Ke's avatar
Guolin Ke committed
14
15
16
17
18
19
20
21
#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
22
 public:
23
24
  DataPartition(data_size_t num_data, int num_leaves)
    :num_data_(num_data), num_leaves_(num_leaves) {
Guolin Ke's avatar
Guolin Ke committed
25
26
27
28
29
    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
30
    used_data_indices_ = nullptr;
Guolin Ke's avatar
Guolin Ke committed
31
32
    #pragma omp parallel
    #pragma omp master
Guolin Ke's avatar
Guolin Ke committed
33
34
35
    {
      num_threads_ = omp_get_num_threads();
    }
Guolin Ke's avatar
Guolin Ke committed
36
37
38
39
40
    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
41
  }
42
43
44
45
46
47

  void ResetLeaves(int num_leaves) {
    num_leaves_ = num_leaves;
    leaf_begin_.resize(num_leaves_);
    leaf_count_.resize(num_leaves_);
  }
48

Guolin Ke's avatar
Guolin Ke committed
49
50
51
52
53
54
  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_);
  }
55

Guolin Ke's avatar
Guolin Ke committed
56
57
58
59
60
61
62
  ~DataPartition() {
  }

  /*!
  * \brief Init, will put all data on the root(leaf_idx = 0)
  */
  void Init() {
63
64
    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
65
66
67
    if (used_data_indices_ == nullptr) {
      // if using all data
      leaf_count_[0] = num_data_;
Guolin Ke's avatar
Guolin Ke committed
68
#pragma omp parallel for schedule(static, 512) if (num_data_ >= 1024)
Guolin Ke's avatar
Guolin Ke committed
69
70
71
72
73
74
      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
75
      std::memcpy(indices_.data(), used_data_indices_, used_data_count_ * sizeof(data_size_t));
Guolin Ke's avatar
Guolin Ke committed
76
77
78
    }
  }

79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
  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
94
95
96
97
98
99
  /*!
  * \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
100
  const data_size_t* GetIndexOnLeaf(int leaf, data_size_t* out_len) const {
Guolin Ke's avatar
Guolin Ke committed
101
102
    // copy reference, maybe unsafe, but faster
    data_size_t begin = leaf_begin_[leaf];
Guolin Ke's avatar
Guolin Ke committed
103
104
    *out_len = leaf_count_[leaf];
    return indices_.data() + begin;
Guolin Ke's avatar
Guolin Ke committed
105
106
107
108
109
110
111
112
113
  }

  /*!
  * \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
  */
114
115
116
117
  void Split(int leaf, const Dataset* dataset, int feature,
             const uint32_t* threshold, int num_threshold, bool default_left,
             int right_leaf) {
    Common::FunctionTimer fun_timer("DataPartition::Split", global_timer);
Guolin Ke's avatar
Guolin Ke committed
118
119
120
121
    // get leaf boundary
    const data_size_t begin = leaf_begin_[leaf];
    const data_size_t cnt = leaf_count_[leaf];

Guolin Ke's avatar
Guolin Ke committed
122
123
124
125
    int nblock = 1;
    data_size_t inner_size = cnt;
    Threading::BlockInfo<data_size_t>(num_threads_, cnt, 512, &nblock,
                                      &inner_size);
126
127
    auto left_start = indices_.data() + begin;
    global_timer.Start("DataPartition::Split.MT");
Guolin Ke's avatar
Guolin Ke committed
128
    // split data multi-threading
129
    OMP_INIT_EX();
Guolin Ke's avatar
Guolin Ke committed
130
#pragma omp parallel for schedule(static, 1)
131
    for (int i = 0; i < nblock; ++i) {
132
      OMP_LOOP_EX_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
133
      data_size_t cur_start = i * inner_size;
134
135
136
137
138
139
      data_size_t cur_cnt = std::min(inner_size, cnt - cur_start);
      if (cur_cnt <= 0) {
        left_cnts_buf_[i] = 0;
        right_cnts_buf_[i] = 0;
        continue;
      }
Guolin Ke's avatar
Guolin Ke committed
140
      // split data inner, reduce the times of function called
Guolin Ke's avatar
Guolin Ke committed
141
142
143
144
145
      data_size_t cur_left_count =
          dataset->Split(feature, threshold, num_threshold, default_left,
                         left_start + cur_start, cur_cnt,
                         temp_left_indices_.data() + cur_start,
                         temp_right_indices_.data() + cur_start);
Guolin Ke's avatar
Guolin Ke committed
146
147
148
      offsets_buf_[i] = cur_start;
      left_cnts_buf_[i] = cur_left_count;
      right_cnts_buf_[i] = cur_cnt - cur_left_count;
149
      OMP_LOOP_EX_END();
Guolin Ke's avatar
Guolin Ke committed
150
    }
151
    OMP_THROW_EX();
152
153
    global_timer.Stop("DataPartition::Split.MT");
    global_timer.Start("DataPartition::Split.Merge");
Guolin Ke's avatar
Guolin Ke committed
154
155
    left_write_pos_buf_[0] = 0;
    right_write_pos_buf_[0] = 0;
156
    for (int i = 1; i < nblock; ++i) {
Guolin Ke's avatar
Guolin Ke committed
157
158
159
160
      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];
Guolin Ke's avatar
Guolin Ke committed
161
    }
Guolin Ke's avatar
Guolin Ke committed
162
163
    data_size_t left_cnt =
        left_write_pos_buf_[nblock - 1] + left_cnts_buf_[nblock - 1];
164
165

    auto right_start = left_start + left_cnt;
Guolin Ke's avatar
Guolin Ke committed
166
#pragma omp parallel for schedule(static)
167
168
169
170
171
    for (int i = 0; i < nblock; ++i) {
      std::copy_n(temp_left_indices_.data() + offsets_buf_[i],
                  left_cnts_buf_[i], left_start + left_write_pos_buf_[i]);
      std::copy_n(temp_right_indices_.data() + offsets_buf_[i],
                  right_cnts_buf_[i], right_start + right_write_pos_buf_[i]);
Guolin Ke's avatar
Guolin Ke committed
172
173
174
175
176
    }
    // update leaf boundary
    leaf_count_[leaf] = left_cnt;
    leaf_begin_[right_leaf] = left_cnt + begin;
    leaf_count_[right_leaf] = cnt - left_cnt;
177
    global_timer.Stop("DataPartition::Split.Merge");
Guolin Ke's avatar
Guolin Ke committed
178
179
180
181
182
183
184
  }

  /*!
  * \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
185
  void SetUsedDataIndices(const data_size_t* used_data_indices, data_size_t num_used_data) {
Guolin Ke's avatar
Guolin Ke committed
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
    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
204
  const data_size_t* indices() const { return indices_.data(); }
Guolin Ke's avatar
Guolin Ke committed
205
206
207
208

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

Nikita Titov's avatar
Nikita Titov committed
209
 private:
Guolin Ke's avatar
Guolin Ke committed
210
211
212
213
214
  /*! \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
215
  std::vector<data_size_t> leaf_begin_;
Guolin Ke's avatar
Guolin Ke committed
216
  /*! \brief number of data on one leaf */
Guolin Ke's avatar
Guolin Ke committed
217
  std::vector<data_size_t> leaf_count_;
Guolin Ke's avatar
Guolin Ke committed
218
  /*! \brief Store all data's indices, order by leaf[data_in_leaf0,..,data_leaf1,..] */
219
  std::vector<data_size_t, Common::AlignmentAllocator<data_size_t, kAlignedSize>> indices_;
Guolin Ke's avatar
Guolin Ke committed
220
  /*! \brief team indices buffer for split */
221
  std::vector<data_size_t, Common::AlignmentAllocator<data_size_t, kAlignedSize>> temp_left_indices_;
Guolin Ke's avatar
Guolin Ke committed
222
  /*! \brief team indices buffer for split */
223
  std::vector<data_size_t, Common::AlignmentAllocator<data_size_t, kAlignedSize>> temp_right_indices_;
Guolin Ke's avatar
Guolin Ke committed
224
225
226
227
228
229
230
  /*! \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
231
  std::vector<data_size_t> offsets_buf_;
Guolin Ke's avatar
Guolin Ke committed
232
  /*! \brief Buffer for multi-threading data partition, used to store left count after split for different threads */
Guolin Ke's avatar
Guolin Ke committed
233
  std::vector<data_size_t> left_cnts_buf_;
Guolin Ke's avatar
Guolin Ke committed
234
  /*! \brief Buffer for multi-threading data partition, used to store right count after split for different threads */
Guolin Ke's avatar
Guolin Ke committed
235
  std::vector<data_size_t> right_cnts_buf_;
Guolin Ke's avatar
Guolin Ke committed
236
  /*! \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
237
  std::vector<data_size_t> left_write_pos_buf_;
Guolin Ke's avatar
Guolin Ke committed
238
  /*! \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
239
  std::vector<data_size_t> right_write_pos_buf_;
Guolin Ke's avatar
Guolin Ke committed
240
241
242
};

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