feature_histogram.hpp 20.4 KB
Newer Older
Guolin Ke's avatar
Guolin Ke committed
1
2
3
4
#ifndef LIGHTGBM_TREELEARNER_FEATURE_HISTOGRAM_HPP_
#define LIGHTGBM_TREELEARNER_FEATURE_HISTOGRAM_HPP_

#include "split_info.hpp"
Guolin Ke's avatar
Guolin Ke committed
5
6
7

#include <LightGBM/utils/array_args.h>
#include <LightGBM/dataset.h>
Guolin Ke's avatar
Guolin Ke committed
8
9
10

#include <cstring>

Guolin Ke's avatar
Guolin Ke committed
11
namespace LightGBM
Guolin Ke's avatar
Guolin Ke committed
12
{
Guolin Ke's avatar
Guolin Ke committed
13

Guolin Ke's avatar
Guolin Ke committed
14
15
16
class FeatureMetainfo {
public:
  int num_bin;
Guolin Ke's avatar
Guolin Ke committed
17
  MissingType missing_type;
18
  int8_t bias = 0;
Guolin Ke's avatar
Guolin Ke committed
19
  uint32_t default_bin;
Guolin Ke's avatar
Guolin Ke committed
20
21
22
  /*! \brief pointer of tree config */
  const TreeConfig* tree_config;
};
Guolin Ke's avatar
Guolin Ke committed
23
24
25
26
27
/*!
* \brief FeatureHistogram is used to construct and store a histogram for a feature.
*/
class FeatureHistogram {
public:
Guolin Ke's avatar
Guolin Ke committed
28
  FeatureHistogram() {
Guolin Ke's avatar
Guolin Ke committed
29
    data_ = nullptr;
Guolin Ke's avatar
Guolin Ke committed
30
  }
Guolin Ke's avatar
Guolin Ke committed
31

Guolin Ke's avatar
Guolin Ke committed
32
33
34
  ~FeatureHistogram() {
  }

Guolin Ke's avatar
Guolin Ke committed
35
36
37
38
39
  /*! \brief Disable copy */
  FeatureHistogram& operator=(const FeatureHistogram&) = delete;
  /*! \brief Disable copy */
  FeatureHistogram(const FeatureHistogram&) = delete;

Guolin Ke's avatar
Guolin Ke committed
40
41
42
43
44
  /*!
  * \brief Init the feature histogram
  * \param feature the feature data for this histogram
  * \param min_num_data_one_leaf minimal number of data in one leaf
  */
45
  void Init(HistogramBinEntry* data, const FeatureMetainfo* meta, BinType bin_type) {
Guolin Ke's avatar
Guolin Ke committed
46
47
    meta_ = meta;
    data_ = data;
48
49
    if (bin_type == BinType::NumericalBin) {
      find_best_threshold_fun_ = std::bind(&FeatureHistogram::FindBestThresholdNumerical, this, std::placeholders::_1
Guolin Ke's avatar
Guolin Ke committed
50
                                           , std::placeholders::_2, std::placeholders::_3, std::placeholders::_4);
51
52
    } else {
      find_best_threshold_fun_ = std::bind(&FeatureHistogram::FindBestThresholdCategorical, this, std::placeholders::_1
Guolin Ke's avatar
Guolin Ke committed
53
                                           , std::placeholders::_2, std::placeholders::_3, std::placeholders::_4);
54
    }
Guolin Ke's avatar
Guolin Ke committed
55
56
  }

Guolin Ke's avatar
Guolin Ke committed
57
58
  HistogramBinEntry* RawData() {
    return data_;
Guolin Ke's avatar
Guolin Ke committed
59
60
61
62
63
64
  }
  /*!
  * \brief Subtract current histograms with other
  * \param other The histogram that want to subtract
  */
  void Subtract(const FeatureHistogram& other) {
Guolin Ke's avatar
Guolin Ke committed
65
    for (int i = 0; i < meta_->num_bin - meta_->bias; ++i) {
Guolin Ke's avatar
Guolin Ke committed
66
67
68
69
70
      data_[i].cnt -= other.data_[i].cnt;
      data_[i].sum_gradients -= other.data_[i].sum_gradients;
      data_[i].sum_hessians -= other.data_[i].sum_hessians;
    }
  }
71

Guolin Ke's avatar
Guolin Ke committed
72
  void FindBestThreshold(double sum_gradient, double sum_hessian, data_size_t num_data,
Guolin Ke's avatar
Guolin Ke committed
73
                         SplitInfo* output) {
Guolin Ke's avatar
Guolin Ke committed
74
    output->default_left = true;
Guolin Ke's avatar
Guolin Ke committed
75
    output->gain = kMinScore;
76
77
78
79
    find_best_threshold_fun_(sum_gradient, sum_hessian + 2 * kEpsilon, num_data, output);
  }

  void FindBestThresholdNumerical(double sum_gradient, double sum_hessian, data_size_t num_data,
Guolin Ke's avatar
Guolin Ke committed
80
                                  SplitInfo* output) {
Guolin Ke's avatar
Guolin Ke committed
81
82

    is_splittable_ = false;
Guolin Ke's avatar
Guolin Ke committed
83
84
    double gain_shift = GetLeafSplitGain(sum_gradient, sum_hessian,
                                         meta_->tree_config->lambda_l1, meta_->tree_config->lambda_l2);
Guolin Ke's avatar
Guolin Ke committed
85
    double min_gain_shift = gain_shift + meta_->tree_config->min_gain_to_split;
Guolin Ke's avatar
Guolin Ke committed
86
87
88
89
90
91
92
93
    if (meta_->num_bin > 2 && meta_->missing_type != MissingType::None) {
      if (meta_->missing_type == MissingType::Zero) {
        FindBestThresholdSequence(sum_gradient, sum_hessian, num_data, min_gain_shift, output, -1, true, false);
        FindBestThresholdSequence(sum_gradient, sum_hessian, num_data, min_gain_shift, output, 1, true, false);
      } else {
        FindBestThresholdSequence(sum_gradient, sum_hessian, num_data, min_gain_shift, output, -1, false, true);
        FindBestThresholdSequence(sum_gradient, sum_hessian, num_data, min_gain_shift, output, 1, false, true);
      }
94
    } else {
Guolin Ke's avatar
Guolin Ke committed
95
96
97
98
99
      FindBestThresholdSequence(sum_gradient, sum_hessian, num_data, min_gain_shift, output, -1, false, false);
      // fix the direction error when only have 2 bins
      if (meta_->missing_type == MissingType::NaN) {
        output->default_left = false;
      }
Guolin Ke's avatar
Guolin Ke committed
100
    }
Guolin Ke's avatar
Guolin Ke committed
101
    output->gain -= min_gain_shift;
Guolin Ke's avatar
Guolin Ke committed
102
  }
103
104

  void FindBestThresholdCategorical(double sum_gradient, double sum_hessian, data_size_t num_data,
Guolin Ke's avatar
Guolin Ke committed
105
                                    SplitInfo* output) {
Guolin Ke's avatar
Guolin Ke committed
106
    output->default_left = false;
107
    double best_gain = kMinScore;
108
109
110
    data_size_t best_left_count = 0;
    double best_sum_left_gradient = 0.0f;
    double best_sum_left_hessian = 0.0f;
Guolin Ke's avatar
Guolin Ke committed
111
112
    double gain_shift = GetLeafSplitGain(sum_gradient, sum_hessian,
                                         meta_->tree_config->lambda_l1, meta_->tree_config->lambda_l2);
113

114
115
    double min_gain_shift = gain_shift + meta_->tree_config->min_gain_to_split;
    is_splittable_ = false;
116
117
118
119
120
121

    uint32_t best_threshold = 0;
    bool is_full_categorical = meta_->missing_type == MissingType::None;

    int used_bin = meta_->num_bin - 1 + is_full_categorical;

122
    // from right to left, and we don't need data in bin0
123
    for (int t = 0; t < used_bin; ++t) {
124
125
      // if data not enough, or sum hessian too small
      if (data_[t].cnt < meta_->tree_config->min_data_in_leaf
Guolin Ke's avatar
Guolin Ke committed
126
          || data_[t].sum_hessians < meta_->tree_config->min_sum_hessian_in_leaf) continue;
127
128
129
130
131
132
133
134
135
136
      data_size_t other_count = num_data - data_[t].cnt;
      // if data not enough
      if (other_count < meta_->tree_config->min_data_in_leaf) continue;

      double sum_other_hessian = sum_hessian - data_[t].sum_hessians - kEpsilon;
      // if sum hessian too small
      if (sum_other_hessian < meta_->tree_config->min_sum_hessian_in_leaf) continue;

      double sum_other_gradient = sum_gradient - data_[t].sum_gradients;
      // current split gain
Guolin Ke's avatar
Guolin Ke committed
137
138
139
140
      double current_gain = GetLeafSplitGain(sum_other_gradient, sum_other_hessian,
                                             meta_->tree_config->lambda_l1, meta_->tree_config->lambda_l2)
        + GetLeafSplitGain(data_[t].sum_gradients, data_[t].sum_hessians + kEpsilon,
                           meta_->tree_config->lambda_l1, meta_->tree_config->lambda_l2);
141
142
143
144
145
146
147
      // gain with split is worse than without split
      if (current_gain <= min_gain_shift) continue;

      // mark to is splittable
      is_splittable_ = true;
      // better split point
      if (current_gain > best_gain) {
148
        best_threshold = static_cast<uint32_t>(t);
149
150
151
        best_sum_left_gradient = data_[t].sum_gradients;
        best_sum_left_hessian = data_[t].sum_hessians + kEpsilon;
        best_left_count = data_[t].cnt;
152
153
154
        best_gain = current_gain;
      }
    }
155

156
157
    if (is_splittable_) {
      // update split information
158
159
160
      output->num_cat_threshold = 1;
      output->cat_threshold.resize(output->num_cat_threshold);
      output->cat_threshold[0] = best_threshold;
Guolin Ke's avatar
Guolin Ke committed
161
162
      output->left_output = CalculateSplittedLeafOutput(best_sum_left_gradient, best_sum_left_hessian,
                                                        meta_->tree_config->lambda_l1, meta_->tree_config->lambda_l2);
163
164
165
166
      output->left_count = best_left_count;
      output->left_sum_gradient = best_sum_left_gradient;
      output->left_sum_hessian = best_sum_left_hessian - kEpsilon;
      output->right_output = CalculateSplittedLeafOutput(sum_gradient - best_sum_left_gradient,
Guolin Ke's avatar
Guolin Ke committed
167
168
                                                         sum_hessian - best_sum_left_hessian,
                                                         meta_->tree_config->lambda_l1, meta_->tree_config->lambda_l2);
169
170
171
      output->right_count = num_data - best_left_count;
      output->right_sum_gradient = sum_gradient - best_sum_left_gradient;
      output->right_sum_hessian = sum_hessian - best_sum_left_hessian - kEpsilon;
Guolin Ke's avatar
Guolin Ke committed
172
      output->gain = best_gain - min_gain_shift;
173
    }
174
175
  }

Guolin Ke's avatar
Guolin Ke committed
176
177
178
179
  /*!
  * \brief Binary size of this histogram
  */
  int SizeOfHistgram() const {
Guolin Ke's avatar
Guolin Ke committed
180
    return (meta_->num_bin - meta_->bias) * sizeof(HistogramBinEntry);
Guolin Ke's avatar
Guolin Ke committed
181
182
183
184
185
  }

  /*!
  * \brief Restore histogram from memory
  */
Guolin Ke's avatar
Guolin Ke committed
186
187
  void FromMemory(char* memory_data) {
    std::memcpy(data_, memory_data, (meta_->num_bin - meta_->bias) * sizeof(HistogramBinEntry));
Guolin Ke's avatar
Guolin Ke committed
188
189
190
191
192
193
194
195
196
197
198
199
200
  }

  /*!
  * \brief True if this histogram can be splitted
  */
  bool is_splittable() { return is_splittable_; }

  /*!
  * \brief Set splittable to this histogram
  */
  void set_is_splittable(bool val) { is_splittable_ = val; }

  /*!
201
  * \brief Calculate the split gain based on regularized sum_gradients and sum_hessians
Guolin Ke's avatar
Guolin Ke committed
202
203
204
205
  * \param sum_gradients
  * \param sum_hessians
  * \return split gain
  */
Guolin Ke's avatar
Guolin Ke committed
206
  static double GetLeafSplitGain(double sum_gradients, double sum_hessians, double l1, double l2) {
207
    double abs_sum_gradients = std::fabs(sum_gradients);
Guolin Ke's avatar
Guolin Ke committed
208
    double reg_abs_sum_gradients = std::max(0.0, abs_sum_gradients - l1);
Guolin Ke's avatar
Guolin Ke committed
209
    return (reg_abs_sum_gradients * reg_abs_sum_gradients)
Guolin Ke's avatar
Guolin Ke committed
210
      / (sum_hessians + l2);
Guolin Ke's avatar
Guolin Ke committed
211

Guolin Ke's avatar
Guolin Ke committed
212
213
214
  }

  /*!
215
  * \brief Calculate the output of a leaf based on regularized sum_gradients and sum_hessians
Guolin Ke's avatar
Guolin Ke committed
216
217
218
219
  * \param sum_gradients
  * \param sum_hessians
  * \return leaf output
  */
Guolin Ke's avatar
Guolin Ke committed
220
  static double CalculateSplittedLeafOutput(double sum_gradients, double sum_hessians, double l1, double l2) {
221
    double abs_sum_gradients = std::fabs(sum_gradients);
Guolin Ke's avatar
Guolin Ke committed
222
    double reg_abs_sum_gradients = std::max(0.0, abs_sum_gradients - l1);
Guolin Ke's avatar
Guolin Ke committed
223
    return -std::copysign(reg_abs_sum_gradients, sum_gradients)
Guolin Ke's avatar
Guolin Ke committed
224
      / (sum_hessians + l2);
Guolin Ke's avatar
Guolin Ke committed
225
  }
Guolin Ke's avatar
Guolin Ke committed
226
227
228

private:

Guolin Ke's avatar
Guolin Ke committed
229
  void FindBestThresholdSequence(double sum_gradient, double sum_hessian, data_size_t num_data, double min_gain_shift,
Guolin Ke's avatar
Guolin Ke committed
230
                                 SplitInfo* output, int dir, bool skip_default_bin, bool use_na_as_missing) {
Guolin Ke's avatar
Guolin Ke committed
231

232
    const int8_t bias = meta_->bias;
Guolin Ke's avatar
Guolin Ke committed
233
234
235
236
237
238
239
240
241
242
243
244
245

    double best_sum_left_gradient = NAN;
    double best_sum_left_hessian = NAN;
    double best_gain = kMinScore;
    data_size_t best_left_count = 0;
    uint32_t best_threshold = static_cast<uint32_t>(meta_->num_bin);

    if (dir == -1) {

      double sum_right_gradient = 0.0f;
      double sum_right_hessian = kEpsilon;
      data_size_t right_count = 0;

Guolin Ke's avatar
Guolin Ke committed
246
      int t = meta_->num_bin - 1 - bias - use_na_as_missing;
Guolin Ke's avatar
Guolin Ke committed
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
      const int t_end = 1 - bias;

      // from right to left, and we don't need data in bin0
      for (; t >= t_end; --t) {

        // need to skip default bin
        if (skip_default_bin && (t + bias) == static_cast<int>(meta_->default_bin)) { continue; }

        sum_right_gradient += data_[t].sum_gradients;
        sum_right_hessian += data_[t].sum_hessians;
        right_count += data_[t].cnt;
        // if data not enough, or sum hessian too small
        if (right_count < meta_->tree_config->min_data_in_leaf
            || sum_right_hessian < meta_->tree_config->min_sum_hessian_in_leaf) continue;
        data_size_t left_count = num_data - right_count;
        // if data not enough
        if (left_count < meta_->tree_config->min_data_in_leaf) break;

        double sum_left_hessian = sum_hessian - sum_right_hessian;
        // if sum hessian too small
        if (sum_left_hessian < meta_->tree_config->min_sum_hessian_in_leaf) break;

        double sum_left_gradient = sum_gradient - sum_right_gradient;
        // current split gain
        double current_gain = GetLeafSplitGain(sum_left_gradient, sum_left_hessian,
                                               meta_->tree_config->lambda_l1, meta_->tree_config->lambda_l2)
          + GetLeafSplitGain(sum_right_gradient, sum_right_hessian,
                             meta_->tree_config->lambda_l1, meta_->tree_config->lambda_l2);
        // gain with split is worse than without split
        if (current_gain <= min_gain_shift) continue;

        // mark to is splittable
        is_splittable_ = true;
        // better split point
        if (current_gain > best_gain) {
          best_left_count = left_count;
          best_sum_left_gradient = sum_left_gradient;
          best_sum_left_hessian = sum_left_hessian;
          // left is <= threshold, right is > threshold.  so this is t-1
          best_threshold = static_cast<uint32_t>(t - 1 + bias);
          best_gain = current_gain;
        }
      }
    } else{
      double sum_left_gradient = 0.0f;
      double sum_left_hessian = kEpsilon;
      data_size_t left_count = 0;

      int t = 0;
      const int t_end = meta_->num_bin - 2 - bias;

Guolin Ke's avatar
Guolin Ke committed
298
299
300
301
302
303
304
305
306
307
308
309
      if (use_na_as_missing && bias == 1) {
        sum_left_gradient = sum_gradient;
        sum_left_hessian = sum_hessian - kEpsilon;
        left_count = num_data;
        for (int i = 0; i < meta_->num_bin - bias; ++i) {
          sum_left_gradient -= data_[i].sum_gradients;
          sum_left_hessian -= data_[i].sum_hessians;
          left_count -= data_[i].cnt;
        }
        t = -1;
      }

Guolin Ke's avatar
Guolin Ke committed
310
311
312
313
      for (; t <= t_end; ++t) {

        // need to skip default bin
        if (skip_default_bin && (t + bias) == static_cast<int>(meta_->default_bin)) { continue; }
Guolin Ke's avatar
Guolin Ke committed
314
315
316
317
318
        if (t >= 0) {
          sum_left_gradient += data_[t].sum_gradients;
          sum_left_hessian += data_[t].sum_hessians;
          left_count += data_[t].cnt;
        }
Guolin Ke's avatar
Guolin Ke committed
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
        // if data not enough, or sum hessian too small
        if (left_count < meta_->tree_config->min_data_in_leaf
            || sum_left_hessian < meta_->tree_config->min_sum_hessian_in_leaf) continue;
        data_size_t right_count = num_data - left_count;
        // if data not enough
        if (right_count < meta_->tree_config->min_data_in_leaf) break;

        double sum_right_hessian = sum_hessian - sum_left_hessian;
        // if sum hessian too small
        if (sum_right_hessian < meta_->tree_config->min_sum_hessian_in_leaf) break;

        double sum_right_gradient = sum_gradient - sum_left_gradient;
        // current split gain
        double current_gain = GetLeafSplitGain(sum_left_gradient, sum_left_hessian,
                                               meta_->tree_config->lambda_l1, meta_->tree_config->lambda_l2)
          + GetLeafSplitGain(sum_right_gradient, sum_right_hessian,
                             meta_->tree_config->lambda_l1, meta_->tree_config->lambda_l2);
        // gain with split is worse than without split
        if (current_gain <= min_gain_shift) continue;

        // mark to is splittable
        is_splittable_ = true;
        // better split point
        if (current_gain > best_gain) {
          best_left_count = left_count;
          best_sum_left_gradient = sum_left_gradient;
          best_sum_left_hessian = sum_left_hessian;
          best_threshold = static_cast<uint32_t>(t + bias);
          best_gain = current_gain;
        }
      }
    }

    if (is_splittable_ && best_gain > output->gain) {
      // update split information
      output->threshold = best_threshold;
      output->left_output = CalculateSplittedLeafOutput(best_sum_left_gradient, best_sum_left_hessian,
                                                        meta_->tree_config->lambda_l1, meta_->tree_config->lambda_l2);
      output->left_count = best_left_count;
      output->left_sum_gradient = best_sum_left_gradient;
      output->left_sum_hessian = best_sum_left_hessian - kEpsilon;
      output->right_output = CalculateSplittedLeafOutput(sum_gradient - best_sum_left_gradient,
                                                         sum_hessian - best_sum_left_hessian,
                                                         meta_->tree_config->lambda_l1, meta_->tree_config->lambda_l2);
      output->right_count = num_data - best_left_count;
      output->right_sum_gradient = sum_gradient - best_sum_left_gradient;
      output->right_sum_hessian = sum_hessian - best_sum_left_hessian - kEpsilon;
      output->gain = best_gain;
Guolin Ke's avatar
Guolin Ke committed
367
      output->default_left = dir == -1;
Guolin Ke's avatar
Guolin Ke committed
368
369
370
    }
  }

Guolin Ke's avatar
Guolin Ke committed
371
  const FeatureMetainfo* meta_;
Guolin Ke's avatar
Guolin Ke committed
372
  /*! \brief sum of gradient of each bin */
Guolin Ke's avatar
Guolin Ke committed
373
374
  HistogramBinEntry* data_;
  //std::vector<HistogramBinEntry> data_;
Guolin Ke's avatar
Guolin Ke committed
375
376
  /*! \brief False if this histogram cannot split */
  bool is_splittable_ = true;
377
378

  std::function<void(double, double, data_size_t, SplitInfo*)> find_best_threshold_fun_;
Guolin Ke's avatar
Guolin Ke committed
379
};
Guolin Ke's avatar
Guolin Ke committed
380
381
382
383
384
385
class HistogramPool {
public:
  /*!
  * \brief Constructor
  */
  HistogramPool() {
Guolin Ke's avatar
Guolin Ke committed
386
387
    cache_size_ = 0;
    total_size_ = 0;
Guolin Ke's avatar
Guolin Ke committed
388
389
390
391
392
393
394
395
396
397
398
  }
  /*!
  * \brief Destructor
  */
  ~HistogramPool() {
  }
  /*!
  * \brief Reset pool size
  * \param cache_size Max cache size
  * \param total_size Total size will be used
  */
Guolin Ke's avatar
Guolin Ke committed
399
  void Reset(int cache_size, int total_size) {
Guolin Ke's avatar
Guolin Ke committed
400
401
402
403
404
405
406
407
408
    cache_size_ = cache_size;
    // at least need 2 bucket to store smaller leaf and larger leaf
    CHECK(cache_size_ >= 2);
    total_size_ = total_size;
    if (cache_size_ > total_size_) {
      cache_size_ = total_size_;
    }
    is_enough_ = (cache_size_ == total_size_);
    if (!is_enough_) {
409
410
411
      mapper_.resize(total_size_);
      inverse_mapper_.resize(cache_size_);
      last_used_time_.resize(cache_size_);
Guolin Ke's avatar
Guolin Ke committed
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
      ResetMap();
    }
  }
  /*!
  * \brief Reset mapper
  */
  void ResetMap() {
    if (!is_enough_) {
      cur_time_ = 0;
      std::fill(mapper_.begin(), mapper_.end(), -1);
      std::fill(inverse_mapper_.begin(), inverse_mapper_.end(), -1);
      std::fill(last_used_time_.begin(), last_used_time_.end(), 0);
    }
  }

Guolin Ke's avatar
Guolin Ke committed
427
428
  void DynamicChangeSize(const Dataset* train_data, const TreeConfig* tree_config, int cache_size, int total_size) {
    if (feature_metas_.empty()) {
Guolin Ke's avatar
Guolin Ke committed
429
430
431
432
      int num_feature = train_data->num_features();
      feature_metas_.resize(num_feature);
      #pragma omp parallel for schedule(static, 512) if(num_feature >= 1024)
      for (int i = 0; i < num_feature; ++i) {
Guolin Ke's avatar
Guolin Ke committed
433
        feature_metas_[i].num_bin = train_data->FeatureNumBin(i);
Guolin Ke's avatar
Guolin Ke committed
434
        feature_metas_[i].default_bin = train_data->FeatureBinMapper(i)->GetDefaultBin();
Guolin Ke's avatar
Guolin Ke committed
435
        feature_metas_[i].missing_type = train_data->FeatureBinMapper(i)->missing_type();
Guolin Ke's avatar
Guolin Ke committed
436
437
438
439
440
441
442
        if (train_data->FeatureBinMapper(i)->GetDefaultBin() == 0) {
          feature_metas_[i].bias = 1;
        } else {
          feature_metas_[i].bias = 0;
        }
        feature_metas_[i].tree_config = tree_config;
      }
Guolin Ke's avatar
Guolin Ke committed
443
    }
Guolin Ke's avatar
Guolin Ke committed
444
445
    uint64_t num_total_bin = train_data->NumTotalBin();
    Log::Info("Total Bins %d", num_total_bin);
Guolin Ke's avatar
Guolin Ke committed
446
    int old_cache_size = static_cast<int>(pool_.size());
Guolin Ke's avatar
Guolin Ke committed
447
    Reset(cache_size, total_size);
Guolin Ke's avatar
Guolin Ke committed
448
449
450
451
452
453

    if (cache_size > old_cache_size) {
      pool_.resize(cache_size);
      data_.resize(cache_size);
    }

454
    OMP_INIT_EX();
Guolin Ke's avatar
Guolin Ke committed
455
    #pragma omp parallel for schedule(static)
Guolin Ke's avatar
Guolin Ke committed
456
    for (int i = old_cache_size; i < cache_size; ++i) {
457
      OMP_LOOP_EX_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
458
459
460
461
462
      pool_[i].reset(new FeatureHistogram[train_data->num_features()]);
      data_[i].resize(num_total_bin);
      uint64_t offset = 0;
      for (int j = 0; j < train_data->num_features(); ++j) {
        offset += static_cast<uint64_t>(train_data->SubFeatureBinOffset(j));
463
        pool_[i][j].Init(data_[i].data() + offset, &feature_metas_[j], train_data->FeatureBinMapper(j)->bin_type());
Guolin Ke's avatar
Guolin Ke committed
464
465
466
467
468
469
470
        auto num_bin = train_data->FeatureNumBin(j);
        if (train_data->FeatureBinMapper(j)->GetDefaultBin() == 0) {
          num_bin -= 1;
        }
        offset += static_cast<uint64_t>(num_bin);
      }
      CHECK(offset == num_total_bin);
471
      OMP_LOOP_EX_END();
Guolin Ke's avatar
Guolin Ke committed
472
    }
473
    OMP_THROW_EX();
Guolin Ke's avatar
Guolin Ke committed
474
475
  }

Guolin Ke's avatar
Guolin Ke committed
476
  void ResetConfig(const TreeConfig* tree_config) {
Guolin Ke's avatar
Guolin Ke committed
477
478
479
    int size = static_cast<int>(feature_metas_.size());
    #pragma omp parallel for schedule(static, 512) if(size >= 1024)
    for (int i = 0; i < size; ++i) {
Guolin Ke's avatar
Guolin Ke committed
480
      feature_metas_[i].tree_config = tree_config;
Guolin Ke's avatar
Guolin Ke committed
481
482
    }
  }
Guolin Ke's avatar
Guolin Ke committed
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
  /*!
  * \brief Get data for the specific index
  * \param idx which index want to get
  * \param out output data will store into this
  * \return True if this index is in the pool, False if this index is not in the pool
  */
  bool Get(int idx, FeatureHistogram** out) {
    if (is_enough_) {
      *out = pool_[idx].get();
      return true;
    } else if (mapper_[idx] >= 0) {
      int slot = mapper_[idx];
      *out = pool_[slot].get();
      last_used_time_[slot] = ++cur_time_;
      return true;
    } else {
      // choose the least used slot 
      int slot = static_cast<int>(ArrayArgs<int>::ArgMin(last_used_time_));
      *out = pool_[slot].get();
      last_used_time_[slot] = ++cur_time_;

      // reset previous mapper
      if (inverse_mapper_[slot] >= 0) mapper_[inverse_mapper_[slot]] = -1;

      // update current mapper
      mapper_[idx] = slot;
      inverse_mapper_[slot] = idx;
      return false;
    }
  }

  /*!
  * \brief Move data from one index to another index
  * \param src_idx
  * \param dst_idx
  */
  void Move(int src_idx, int dst_idx) {
    if (is_enough_) {
      std::swap(pool_[src_idx], pool_[dst_idx]);
      return;
    }
    if (mapper_[src_idx] < 0) {
      return;
    }
    // get slot of src idx
    int slot = mapper_[src_idx];
    // reset src_idx
    mapper_[src_idx] = -1;

    // move to dst idx
    mapper_[dst_idx] = slot;
    last_used_time_[slot] = ++cur_time_;
    inverse_mapper_[slot] = dst_idx;
  }
private:
  std::vector<std::unique_ptr<FeatureHistogram[]>> pool_;
Guolin Ke's avatar
Guolin Ke committed
539
540
  std::vector<std::vector<HistogramBinEntry>> data_;
  std::vector<FeatureMetainfo> feature_metas_;
Guolin Ke's avatar
Guolin Ke committed
541
542
543
544
545
546
547
548
549
  int cache_size_;
  int total_size_;
  bool is_enough_ = false;
  std::vector<int> mapper_;
  std::vector<int> inverse_mapper_;
  std::vector<int> last_used_time_;
  int cur_time_ = 0;
};

Guolin Ke's avatar
Guolin Ke committed
550
}  // namespace LightGBM
Guolin Ke's avatar
Guolin Ke committed
551
#endif   // LightGBM_TREELEARNER_FEATURE_HISTOGRAM_HPP_