tree.cpp 15.8 KB
Newer Older
Guolin Ke's avatar
Guolin Ke committed
1
2
3
4
5
6
7
8
9
10
11
12
#include <LightGBM/tree.h>

#include <LightGBM/utils/threading.h>
#include <LightGBM/utils/common.h>

#include <LightGBM/dataset.h>

#include <sstream>
#include <unordered_map>
#include <functional>
#include <vector>
#include <string>
Guolin Ke's avatar
Guolin Ke committed
13
#include <memory>
14
#include <iomanip>
Guolin Ke's avatar
Guolin Ke committed
15
16
17
18
19
20

namespace LightGBM {

Tree::Tree(int max_leaves)
  :max_leaves_(max_leaves) {

Guolin Ke's avatar
Guolin Ke committed
21
22
23
24
25
26
  left_child_.resize(max_leaves_ - 1);
  right_child_.resize(max_leaves_ - 1);
  split_feature_inner_.resize(max_leaves_ - 1);
  split_feature_.resize(max_leaves_ - 1);
  threshold_in_bin_.resize(max_leaves_ - 1);
  threshold_.resize(max_leaves_ - 1);
Guolin Ke's avatar
Guolin Ke committed
27
  decision_type_.resize(max_leaves_ - 1, 0);
Guolin Ke's avatar
Guolin Ke committed
28
29
30
31
32
33
34
  split_gain_.resize(max_leaves_ - 1);
  leaf_parent_.resize(max_leaves_);
  leaf_value_.resize(max_leaves_);
  leaf_count_.resize(max_leaves_);
  internal_value_.resize(max_leaves_ - 1);
  internal_count_.resize(max_leaves_ - 1);
  leaf_depth_.resize(max_leaves_);
Guolin Ke's avatar
Guolin Ke committed
35
36
  // root is in the depth 0
  leaf_depth_[0] = 0;
Guolin Ke's avatar
Guolin Ke committed
37
  num_leaves_ = 1;
38
  leaf_value_[0] = 0.0f;
Guolin Ke's avatar
Guolin Ke committed
39
  leaf_parent_[0] = -1;
Guolin Ke's avatar
Guolin Ke committed
40
  shrinkage_ = 1.0f;
41
  num_cat_ = 0;
Guolin Ke's avatar
Guolin Ke committed
42
}
Guolin Ke's avatar
Guolin Ke committed
43

Guolin Ke's avatar
Guolin Ke committed
44
Tree::~Tree() {
Guolin Ke's avatar
Guolin Ke committed
45

Guolin Ke's avatar
Guolin Ke committed
46
47
}

48
49
50
51
int Tree::Split(int leaf, int feature, int real_feature, uint32_t threshold_bin,
                double threshold_double, double left_value, double right_value,
                data_size_t left_cnt, data_size_t right_cnt, double gain, MissingType missing_type, bool default_left) {
  Split(leaf, feature, real_feature, left_value, right_value, left_cnt, right_cnt, gain);
Guolin Ke's avatar
Guolin Ke committed
52
  int new_node_idx = num_leaves_ - 1;
Guolin Ke's avatar
Guolin Ke committed
53
  decision_type_[new_node_idx] = 0;
54
  SetDecisionType(&decision_type_[new_node_idx], false, kCategoricalMask);
Guolin Ke's avatar
Guolin Ke committed
55
56
57
58
59
60
61
  SetDecisionType(&decision_type_[new_node_idx], default_left, kDefaultLeftMask);
  if (missing_type == MissingType::None) {
    SetMissingType(&decision_type_[new_node_idx], 0);
  } else if (missing_type == MissingType::Zero) {
    SetMissingType(&decision_type_[new_node_idx], 1);
  } else if (missing_type == MissingType::NaN) {
    SetMissingType(&decision_type_[new_node_idx], 2);
62
  }
Guolin Ke's avatar
Guolin Ke committed
63
  threshold_in_bin_[new_node_idx] = threshold_bin;
64
  threshold_[new_node_idx] = Common::AvoidInf(threshold_double);
65
66
67
  ++num_leaves_;
  return num_leaves_ - 1;
}
Guolin Ke's avatar
Guolin Ke committed
68

69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
int Tree::SplitCategorical(int leaf, int feature, int real_feature, uint32_t threshold_bin,
                           double threshold, double left_value, double right_value,
                           data_size_t left_cnt, data_size_t right_cnt, double gain, MissingType missing_type) {
  Split(leaf, feature, real_feature, left_value, right_value, left_cnt, right_cnt, gain);
  int new_node_idx = num_leaves_ - 1;
  decision_type_[new_node_idx] = 0;
  SetDecisionType(&decision_type_[new_node_idx], true, kCategoricalMask);
  if (missing_type == MissingType::None) {
    SetMissingType(&decision_type_[new_node_idx], 0);
  } else if (missing_type == MissingType::Zero) {
    SetMissingType(&decision_type_[new_node_idx], 1);
  } else if (missing_type == MissingType::NaN) {
    SetMissingType(&decision_type_[new_node_idx], 2);
  }
  threshold_in_bin_[new_node_idx] = threshold_bin;
  threshold_[new_node_idx] = threshold;
  ++num_cat_;
Guolin Ke's avatar
Guolin Ke committed
86
87
88
89
  ++num_leaves_;
  return num_leaves_ - 1;
}

90
91
92
93
94
95
96
97
98
99
100
101
102
103
#define PredictionFun(niter, fidx_in_iter, start_pos, decision_fun, iter_idx, data_idx) \
std::vector<std::unique_ptr<BinIterator>> iter((niter)); \
for (int i = 0; i < (niter); ++i) { \
  iter[i].reset(data->FeatureIterator((fidx_in_iter))); \
  iter[i]->Reset((start_pos)); \
}\
for (data_size_t i = start; i < end; ++i) {\
  int node = 0;\
  while (node >= 0) {\
    node = decision_fun(iter[(iter_idx)]->Get((data_idx)), node, default_bins[node], max_bins[node]);\
  }\
  score[(data_idx)] += static_cast<double>(leaf_value_[~node]);\
}\

104
void Tree::AddPredictionToScore(const Dataset* data, data_size_t num_data, double* score) const {
Guolin Ke's avatar
Guolin Ke committed
105
  if (num_leaves_ <= 1) { return; }
Guolin Ke's avatar
Guolin Ke committed
106
107
108
109
110
111
112
113
  std::vector<uint32_t> default_bins(num_leaves_ - 1);
  std::vector<uint32_t> max_bins(num_leaves_ - 1);
  for (int i = 0; i < num_leaves_ - 1; ++i) {
    const int fidx = split_feature_inner_[i];
    auto bin_mapper = data->FeatureBinMapper(fidx);
    default_bins[i] = bin_mapper->GetDefaultBin();
    max_bins[i] = bin_mapper->num_bin() - 1;
  }
114
  if (num_cat_ > 0) {
115
    if (data->num_features() > num_leaves_ - 1) {
116
117
118
      Threading::For<data_size_t>(0, num_data, [this, &data, score, &default_bins, &max_bins]
      (int, data_size_t start, data_size_t end) {
        PredictionFun(num_leaves_ - 1, split_feature_inner_[i], start, DecisionInner, node, i);
119
120
      });
    } else {
121
122
123
      Threading::For<data_size_t>(0, num_data, [this, &data, score, &default_bins, &max_bins]
      (int, data_size_t start, data_size_t end) {
        PredictionFun(data->num_features(), i, start, DecisionInner, split_feature_inner_[node], i);
124
125
      });
    }
Guolin Ke's avatar
Guolin Ke committed
126
  } else {
127
    if (data->num_features() > num_leaves_ - 1) {
128
129
130
      Threading::For<data_size_t>(0, num_data, [this, &data, score, &default_bins, &max_bins]
      (int, data_size_t start, data_size_t end) {
        PredictionFun(num_leaves_ - 1, split_feature_inner_[i], start, NumericalDecisionInner, node, i);
131
132
      });
    } else {
133
134
135
      Threading::For<data_size_t>(0, num_data, [this, &data, score, &default_bins, &max_bins]
      (int, data_size_t start, data_size_t end) {
        PredictionFun(data->num_features(), i, start, NumericalDecisionInner, split_feature_inner_[node], i);
136
137
      });
    }
Guolin Ke's avatar
Guolin Ke committed
138
  }
Guolin Ke's avatar
Guolin Ke committed
139
140
}

Guolin Ke's avatar
Guolin Ke committed
141
142
143
void Tree::AddPredictionToScore(const Dataset* data,
  const data_size_t* used_data_indices,
  data_size_t num_data, double* score) const {
Guolin Ke's avatar
Guolin Ke committed
144
  if (num_leaves_ <= 1) { return; }
Guolin Ke's avatar
Guolin Ke committed
145
146
147
148
149
150
151
152
  std::vector<uint32_t> default_bins(num_leaves_ - 1);
  std::vector<uint32_t> max_bins(num_leaves_ - 1);
  for (int i = 0; i < num_leaves_ - 1; ++i) {
    const int fidx = split_feature_inner_[i];
    auto bin_mapper = data->FeatureBinMapper(fidx);
    default_bins[i] = bin_mapper->GetDefaultBin();
    max_bins[i] = bin_mapper->num_bin() - 1;
  }
153
  if (num_cat_ > 0) {
154
    if (data->num_features() > num_leaves_ - 1) {
155
156
157
      Threading::For<data_size_t>(0, num_data, [this, &data, score, used_data_indices, &default_bins, &max_bins]
      (int, data_size_t start, data_size_t end) {
        PredictionFun(num_leaves_ - 1, split_feature_inner_[i], used_data_indices[start], DecisionInner, node, used_data_indices[i]);
158
159
      });
    } else {
160
161
162
      Threading::For<data_size_t>(0, num_data, [this, &data, score, used_data_indices, &default_bins, &max_bins]
      (int, data_size_t start, data_size_t end) {
        PredictionFun(data->num_features(), i, used_data_indices[start], DecisionInner, split_feature_inner_[node], used_data_indices[i]);
163
164
      });
    }
Guolin Ke's avatar
Guolin Ke committed
165
  } else {
166
    if (data->num_features() > num_leaves_ - 1) {
167
168
169
      Threading::For<data_size_t>(0, num_data, [this, &data, score, used_data_indices, &default_bins, &max_bins]
      (int, data_size_t start, data_size_t end) {
        PredictionFun(num_leaves_ - 1, split_feature_inner_[i], used_data_indices[start], NumericalDecisionInner, node, used_data_indices[i]);
170
171
      });
    } else {
172
173
174
      Threading::For<data_size_t>(0, num_data, [this, &data, score, used_data_indices, &default_bins, &max_bins]
      (int, data_size_t start, data_size_t end) {
        PredictionFun(data->num_features(), i, used_data_indices[start], NumericalDecisionInner, split_feature_inner_[node], used_data_indices[i]);
175
176
      });
    }
Guolin Ke's avatar
Guolin Ke committed
177
  }
Guolin Ke's avatar
Guolin Ke committed
178
179
}

180
181
#undef PredictionFun

Guolin Ke's avatar
Guolin Ke committed
182
std::string Tree::ToString() {
183
184
  std::stringstream str_buf;
  str_buf << "num_leaves=" << num_leaves_ << std::endl;
185
  str_buf << "num_cat=" << num_cat_ << std::endl;
186
  str_buf << "split_feature="
Guolin Ke's avatar
Guolin Ke committed
187
    << Common::ArrayToString<int>(split_feature_, num_leaves_ - 1, ' ') << std::endl;
188
  str_buf << "split_gain="
Guolin Ke's avatar
Guolin Ke committed
189
    << Common::ArrayToString<double>(split_gain_, num_leaves_ - 1, ' ') << std::endl;
190
  str_buf << "threshold="
Guolin Ke's avatar
Guolin Ke committed
191
    << Common::ArrayToString<double>(threshold_, num_leaves_ - 1, ' ') << std::endl;
192
193
  str_buf << "decision_type="
    << Common::ArrayToString<int>(Common::ArrayCast<int8_t, int>(decision_type_), num_leaves_ - 1, ' ') << std::endl;
194
  str_buf << "left_child="
Guolin Ke's avatar
Guolin Ke committed
195
    << Common::ArrayToString<int>(left_child_, num_leaves_ - 1, ' ') << std::endl;
196
  str_buf << "right_child="
Guolin Ke's avatar
Guolin Ke committed
197
    << Common::ArrayToString<int>(right_child_, num_leaves_ - 1, ' ') << std::endl;
198
  str_buf << "leaf_value="
Guolin Ke's avatar
Guolin Ke committed
199
    << Common::ArrayToString<double>(leaf_value_, num_leaves_, ' ') << std::endl;
200
  str_buf << "leaf_count="
Guolin Ke's avatar
Guolin Ke committed
201
    << Common::ArrayToString<data_size_t>(leaf_count_, num_leaves_, ' ') << std::endl;
202
  str_buf << "internal_value="
Guolin Ke's avatar
Guolin Ke committed
203
    << Common::ArrayToString<double>(internal_value_, num_leaves_ - 1, ' ') << std::endl;
204
  str_buf << "internal_count="
Guolin Ke's avatar
Guolin Ke committed
205
    << Common::ArrayToString<data_size_t>(internal_count_, num_leaves_ - 1, ' ') << std::endl;
Guolin Ke's avatar
Guolin Ke committed
206
  str_buf << "shrinkage=" << shrinkage_ << std::endl;
207
208
  str_buf << std::endl;
  return str_buf.str();
Guolin Ke's avatar
Guolin Ke committed
209
210
}

wxchan's avatar
wxchan committed
211
std::string Tree::ToJSON() {
212
  std::stringstream str_buf;
213
  str_buf << std::setprecision(std::numeric_limits<double>::digits10 + 2);
214
  str_buf << "\"num_leaves\":" << num_leaves_ << "," << std::endl;
215
  str_buf << "\"num_cat\":" << num_cat_ << "," << std::endl;
Guolin Ke's avatar
Guolin Ke committed
216
  str_buf << "\"shrinkage\":" << shrinkage_ << "," << std::endl;
wxchan's avatar
wxchan committed
217
  if (num_leaves_ == 1) {
218
    str_buf << "\"tree_structure\":{" << "\"leaf_value\":" << leaf_value_[0] << "}"  << std::endl;
wxchan's avatar
wxchan committed
219
220
221
  } else {
    str_buf << "\"tree_structure\":" << NodeToJSON(0) << std::endl;
  }
wxchan's avatar
wxchan committed
222

223
  return str_buf.str();
wxchan's avatar
wxchan committed
224
225
226
}

std::string Tree::NodeToJSON(int index) {
227
  std::stringstream str_buf;
228
  str_buf << std::setprecision(std::numeric_limits<double>::digits10 + 2);
wxchan's avatar
wxchan committed
229
230
  if (index >= 0) {
    // non-leaf
231
232
    str_buf << "{" << std::endl;
    str_buf << "\"split_index\":" << index << "," << std::endl;
Guolin Ke's avatar
Guolin Ke committed
233
    str_buf << "\"split_feature\":" << split_feature_[index] << "," << std::endl;
234
    str_buf << "\"split_gain\":" << split_gain_[index] << "," << std::endl;
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
    if (GetDecisionType(decision_type_[index], kCategoricalMask)) {
      str_buf << "\"threshold\":" << static_cast<int>(threshold_[index]) << "," << std::endl;
      str_buf << "\"decision_type\":\"==\"," << std::endl;
    } else {
      str_buf << "\"threshold\":" << Common::AvoidInf(threshold_[index]) << "," << std::endl;
      str_buf << "\"decision_type\":\"<=\"," << std::endl;
    }
    if (GetDecisionType(decision_type_[index], kDefaultLeftMask)) {
      str_buf << "\"default_left\":true," << std::endl;
    } else {
      str_buf << "\"default_left\":false," << std::endl;
    }
    uint8_t missing_type = GetMissingType(decision_type_[index]);
    if (missing_type == 0) {
      str_buf << "\"missing_type\":\"None\"," << std::endl;
    } else if (missing_type == 1) {
      str_buf << "\"missing_type\":\"Zero\"," << std::endl;
    } else {
      str_buf << "\"missing_type\":\"NaN\"," << std::endl;
    }
255
256
257
258
259
    str_buf << "\"internal_value\":" << internal_value_[index] << "," << std::endl;
    str_buf << "\"internal_count\":" << internal_count_[index] << "," << std::endl;
    str_buf << "\"left_child\":" << NodeToJSON(left_child_[index]) << "," << std::endl;
    str_buf << "\"right_child\":" << NodeToJSON(right_child_[index]) << std::endl;
    str_buf << "}";
wxchan's avatar
wxchan committed
260
261
262
  } else {
    // leaf
    index = ~index;
263
264
265
266
267
    str_buf << "{" << std::endl;
    str_buf << "\"leaf_index\":" << index << "," << std::endl;
    str_buf << "\"leaf_value\":" << leaf_value_[index] << "," << std::endl;
    str_buf << "\"leaf_count\":" << leaf_count_[index] << std::endl;
    str_buf << "}";
wxchan's avatar
wxchan committed
268
269
  }

270
  return str_buf.str();
wxchan's avatar
wxchan committed
271
272
}

273
274
275
276
277
278
279
std::string Tree::ToIfElse(int index, bool is_predict_leaf_index) {
  std::stringstream str_buf;
  str_buf << "double PredictTree" << index;
  if (is_predict_leaf_index) {
    str_buf << "Leaf";
  }
  str_buf << "(const double* arr) { ";
280
281
  if (num_leaves_ <= 1) {
    str_buf << "return " << leaf_value_[0] << ";";
282
  } else {
283
284
285
286
287
    // use this for the missing value conversion
    str_buf << "double fval = 0.0f; ";
    if (num_cat_ > 0) {
      str_buf << "int int_fval = 0; ";
    }
288
289
290
291
292
293
294
295
296
297
298
    str_buf << NodeToIfElse(0, is_predict_leaf_index);
  }
  str_buf << " }" << std::endl;
  return str_buf.str();
}

std::string Tree::NodeToIfElse(int index, bool is_predict_leaf_index) {
  std::stringstream str_buf;
  str_buf << std::setprecision(std::numeric_limits<double>::digits10 + 2);
  if (index >= 0) {
    // non-leaf
299
    str_buf << "fval = arr[" << split_feature_[index] << "];";
Guolin Ke's avatar
Guolin Ke committed
300
    if (GetDecisionType(decision_type_[index], kCategoricalMask) == 0) {
301
      str_buf << NumericalDecisionIfElse(index);
302
    } else {
303
      str_buf << CategoricalDecisionIfElse(index);
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
    }
    // left subtree
    str_buf << NodeToIfElse(left_child_[index], is_predict_leaf_index);
    str_buf << " } else { ";
    // right subtree
    str_buf << NodeToIfElse(right_child_[index], is_predict_leaf_index);
    str_buf << " }";
  } else {
    // leaf
    str_buf << "return ";
    if (is_predict_leaf_index) {
      str_buf << ~index;
    } else {
      str_buf << leaf_value_[~index];
    }
    str_buf << ";";
  }

  return str_buf.str();
}

Guolin Ke's avatar
Guolin Ke committed
325
Tree::Tree(const std::string& str) {
Guolin Ke's avatar
Guolin Ke committed
326
  std::vector<std::string> lines = Common::SplitLines(str.c_str());
Guolin Ke's avatar
Guolin Ke committed
327
328
329
330
331
332
333
334
335
336
337
  std::unordered_map<std::string, std::string> key_vals;
  for (const std::string& line : lines) {
    std::vector<std::string> tmp_strs = Common::Split(line.c_str(), '=');
    if (tmp_strs.size() == 2) {
      std::string key = Common::Trim(tmp_strs[0]);
      std::string val = Common::Trim(tmp_strs[1]);
      if (key.size() > 0 && val.size() > 0) {
        key_vals[key] = val;
      }
    }
  }
338
  if (key_vals.count("num_leaves") <= 0) {
Guolin Ke's avatar
Guolin Ke committed
339
    Log::Fatal("Tree model should contain num_leaves field.");
Guolin Ke's avatar
Guolin Ke committed
340
341
342
343
  }

  Common::Atoi(key_vals["num_leaves"].c_str(), &num_leaves_);

344
345
346
347
348
349
  if (key_vals.count("num_cat") <= 0) {
    Log::Fatal("Tree model should contain num_cat field.");
  }

  Common::Atoi(key_vals["num_cat"].c_str(), &num_cat_);

350
351
352
353
354
355
  if (key_vals.count("leaf_value")) {
    leaf_value_ = Common::StringToArray<double>(key_vals["leaf_value"], ' ', num_leaves_);
  } else {
    Log::Fatal("Tree model string format error, should contain leaf_value field");
  }

356
357
  if (num_leaves_ <= 1) { return; }

Guolin Ke's avatar
Guolin Ke committed
358
359
360
361
  if (key_vals.count("left_child")) {
    left_child_ = Common::StringToArray<int>(key_vals["left_child"], ' ', num_leaves_ - 1);
  } else {
    Log::Fatal("Tree model string format error, should contain left_child field");
362
363
  }

Guolin Ke's avatar
Guolin Ke committed
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
  if (key_vals.count("right_child")) {
    right_child_ = Common::StringToArray<int>(key_vals["right_child"], ' ', num_leaves_ - 1);
  } else {
    Log::Fatal("Tree model string format error, should contain right_child field");
  }

  if (key_vals.count("split_feature")) {
    split_feature_ = Common::StringToArray<int>(key_vals["split_feature"], ' ', num_leaves_ - 1);
  } else {
    Log::Fatal("Tree model string format error, should contain split_feature field");
  }

  if (key_vals.count("threshold")) {
    threshold_ = Common::StringToArray<double>(key_vals["threshold"], ' ', num_leaves_ - 1);
  } else {
    Log::Fatal("Tree model string format error, should contain threshold field");
  }

  if (key_vals.count("split_gain")) {
    split_gain_ = Common::StringToArray<double>(key_vals["split_gain"], ' ', num_leaves_ - 1);
  } else {
    split_gain_.resize(num_leaves_ - 1);
  }

  if (key_vals.count("internal_count")) {
    internal_count_ = Common::StringToArray<data_size_t>(key_vals["internal_count"], ' ', num_leaves_ - 1);
  } else {
    internal_count_.resize(num_leaves_ - 1);
  }

  if (key_vals.count("internal_value")) {
    internal_value_ = Common::StringToArray<double>(key_vals["internal_value"], ' ', num_leaves_ - 1);
  } else {
    internal_value_.resize(num_leaves_ - 1);
  }

  if (key_vals.count("leaf_count")) {
    leaf_count_ = Common::StringToArray<data_size_t>(key_vals["leaf_count"], ' ', num_leaves_);
  } else {
    leaf_count_.resize(num_leaves_);
  }

  if (key_vals.count("decision_type")) {
    decision_type_ = Common::StringToArray<int8_t>(key_vals["decision_type"], ' ', num_leaves_ - 1);
  } else {
    decision_type_ = std::vector<int8_t>(num_leaves_ - 1, 0);
  }

  if (key_vals.count("shrinkage")) {
    Common::Atof(key_vals["shrinkage"].c_str(), &shrinkage_);
  } else {
    shrinkage_ = 1.0f;
  }
Guolin Ke's avatar
Guolin Ke committed
417
418
419
}

}  // namespace LightGBM