gbdt.cpp 22.1 KB
Newer Older
Guolin Ke's avatar
Guolin Ke committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
#include "gbdt.h"

#include <LightGBM/utils/common.h>

#include <LightGBM/feature.h>
#include <LightGBM/objective_function.h>
#include <LightGBM/metric.h>

#include <ctime>

#include <sstream>
#include <chrono>
#include <string>
#include <vector>
15
#include <utility>
Guolin Ke's avatar
Guolin Ke committed
16
17
18

namespace LightGBM {

wxchan's avatar
wxchan committed
19
20
21
GBDT::GBDT() 
  :num_iteration_for_pred_(0), 
  num_init_iteration_(0) {
Guolin Ke's avatar
Guolin Ke committed
22

Guolin Ke's avatar
Guolin Ke committed
23
24
25
}

GBDT::~GBDT() {
Guolin Ke's avatar
Guolin Ke committed
26

Guolin Ke's avatar
Guolin Ke committed
27
28
}

29
30
31
void GBDT::Init(const BoostingConfig* config, const Dataset* train_data, const ObjectiveFunction* object_function,
     const std::vector<const Metric*>& training_metrics) {
  iter_ = 0;
wxchan's avatar
wxchan committed
32
  num_iteration_for_pred_ = 0;
33
  max_feature_idx_ = 0;
wxchan's avatar
wxchan committed
34
  num_class_ = config->num_class;
Guolin Ke's avatar
Guolin Ke committed
35
  random_ = Random(config->bagging_seed);
wxchan's avatar
wxchan committed
36
  train_data_ = nullptr;
Guolin Ke's avatar
Guolin Ke committed
37
  gbdt_config_ = nullptr;
Guolin Ke's avatar
Guolin Ke committed
38
  tree_learner_.clear();
wxchan's avatar
wxchan committed
39
40
41
42
43
  ResetTrainingData(config, train_data, object_function, training_metrics);
}

void GBDT::ResetTrainingData(const BoostingConfig* config, const Dataset* train_data, const ObjectiveFunction* object_function,
  const std::vector<const Metric*>& training_metrics) {
Guolin Ke's avatar
Guolin Ke committed
44
  auto new_config = std::unique_ptr<BoostingConfig>(new BoostingConfig(*config));
wxchan's avatar
wxchan committed
45
46
47
  if (train_data_ != nullptr && !train_data_->CheckAlign(*train_data)) {
    Log::Fatal("cannot reset training data, since new training data has different bin mappers");
  }
Guolin Ke's avatar
Guolin Ke committed
48
49
50
  early_stopping_round_ = new_config->early_stopping_round;
  shrinkage_rate_ = new_config->learning_rate;

Guolin Ke's avatar
Guolin Ke committed
51
  object_function_ = object_function;
Guolin Ke's avatar
Guolin Ke committed
52

Guolin Ke's avatar
Guolin Ke committed
53
  sigmoid_ = -1.0f;
wxchan's avatar
wxchan committed
54
  if (object_function_ != nullptr
Guolin Ke's avatar
Guolin Ke committed
55
56
    && std::string(object_function_->GetName()) == std::string("binary")) {
    // only binary classification need sigmoid transform
Guolin Ke's avatar
Guolin Ke committed
57
    sigmoid_ = new_config->sigmoid;
58
  }
Guolin Ke's avatar
Guolin Ke committed
59

Guolin Ke's avatar
Guolin Ke committed
60
61
62
63
64
65
66
67
68
69
70
71
72
  if (train_data_ != train_data && train_data != nullptr) {
    if (tree_learner_.empty()) {
      for (int i = 0; i < num_class_; ++i) {
        auto new_tree_learner = std::unique_ptr<TreeLearner>(TreeLearner::CreateTreeLearner(new_config->tree_learner_type, &new_config->tree_config));
        tree_learner_.push_back(std::move(new_tree_learner));
      }
      tree_learner_.shrink_to_fit();
    }
    // init tree learner
    for (int i = 0; i < num_class_; ++i) {
      tree_learner_[i]->Init(train_data);
    }

Guolin Ke's avatar
Guolin Ke committed
73
74
75
76
77
78
    // push training metrics
    training_metrics_.clear();
    for (const auto& metric : training_metrics) {
      training_metrics_.push_back(metric);
    }
    training_metrics_.shrink_to_fit();
wxchan's avatar
wxchan committed
79
80
81
82
83
84
85
86
87
88
89
90
91
    // not same training data, need reset score and others
    // create score tracker
    train_score_updater_.reset(new ScoreUpdater(train_data, num_class_));
    // update score
    for (int i = 0; i < iter_; ++i) {
      for (int curr_class = 0; curr_class < num_class_; ++curr_class) {
        auto curr_tree = (i + num_init_iteration_) * num_class_ + curr_class;
        train_score_updater_->AddScore(models_[curr_tree].get(), curr_class);
      }
    }
    num_data_ = train_data->num_data();
    // create buffer for gradients and hessians
    if (object_function_ != nullptr) {
92
93
      gradients_.resize(num_data_ * num_class_);
      hessians_.resize(num_data_ * num_class_);
wxchan's avatar
wxchan committed
94
95
96
97
98
    }
    // get max feature index
    max_feature_idx_ = train_data->num_total_features() - 1;
    // get label index
    label_idx_ = train_data->label_idx();
Guolin Ke's avatar
Guolin Ke committed
99
100
  }

Guolin Ke's avatar
Guolin Ke committed
101
102
  if ((train_data_ != train_data && train_data != nullptr)
    || (gbdt_config_ != nullptr && gbdt_config_->bagging_fraction != new_config->bagging_fraction)) {
wxchan's avatar
wxchan committed
103
    // if need bagging, create buffer
Guolin Ke's avatar
Guolin Ke committed
104
    if (new_config->bagging_fraction < 1.0 && new_config->bagging_freq > 0) {
105
106
      out_of_bag_data_indices_.resize(num_data_);
      bag_data_indices_.resize(num_data_);
wxchan's avatar
wxchan committed
107
108
109
110
111
112
    } else {
      out_of_bag_data_cnt_ = 0;
      out_of_bag_data_indices_.clear();
      bag_data_cnt_ = num_data_;
      bag_data_indices_.clear();
    }
Guolin Ke's avatar
Guolin Ke committed
113
  }
wxchan's avatar
wxchan committed
114
  train_data_ = train_data;
Guolin Ke's avatar
Guolin Ke committed
115
116
117
118
119
120
  if (train_data_ != nullptr) {
    // reset config for tree learner
    for (int i = 0; i < num_class_; ++i) {
      tree_learner_[i]->ResetConfig(&new_config->tree_config);
    }
  }
Guolin Ke's avatar
Guolin Ke committed
121
  gbdt_config_.reset(new_config.release());
Guolin Ke's avatar
Guolin Ke committed
122
123
}

wxchan's avatar
wxchan committed
124
void GBDT::AddValidDataset(const Dataset* valid_data,
Guolin Ke's avatar
Guolin Ke committed
125
  const std::vector<const Metric*>& valid_metrics) {
wxchan's avatar
wxchan committed
126
127
  if (!train_data_->CheckAlign(*valid_data)) {
    Log::Fatal("cannot add validation data, since it has different bin mappers with training data");
128
  }
Guolin Ke's avatar
Guolin Ke committed
129
  // for a validation dataset, we need its score and metric
Guolin Ke's avatar
Guolin Ke committed
130
  auto new_score_updater = std::unique_ptr<ScoreUpdater>(new ScoreUpdater(valid_data, num_class_));
wxchan's avatar
wxchan committed
131
132
133
134
135
136
137
  // update score
  for (int i = 0; i < iter_; ++i) {
    for (int curr_class = 0; curr_class < num_class_; ++curr_class) {
      auto curr_tree = (i + num_init_iteration_) * num_class_ + curr_class;
      new_score_updater->AddScore(models_[curr_tree].get(), curr_class);
    }
  }
Guolin Ke's avatar
Guolin Ke committed
138
  valid_score_updater_.push_back(std::move(new_score_updater));
Guolin Ke's avatar
Guolin Ke committed
139
  valid_metrics_.emplace_back();
140
141
142
  if (early_stopping_round_ > 0) {
    best_iter_.emplace_back();
    best_score_.emplace_back();
Guolin Ke's avatar
Guolin Ke committed
143
    best_msg_.emplace_back();
144
  }
Guolin Ke's avatar
Guolin Ke committed
145
146
  for (const auto& metric : valid_metrics) {
    valid_metrics_.back().push_back(metric);
147
148
149
    if (early_stopping_round_ > 0) {
      best_iter_.back().push_back(0);
      best_score_.back().push_back(kMinScore);
Guolin Ke's avatar
Guolin Ke committed
150
      best_msg_.back().emplace_back();
151
    }
Guolin Ke's avatar
Guolin Ke committed
152
  }
Guolin Ke's avatar
Guolin Ke committed
153
  valid_metrics_.back().shrink_to_fit();
Guolin Ke's avatar
Guolin Ke committed
154
155
156
}


157
void GBDT::Bagging(int iter, const int curr_class) {
Guolin Ke's avatar
Guolin Ke committed
158
  // if need bagging
Guolin Ke's avatar
Guolin Ke committed
159
  if (!out_of_bag_data_indices_.empty() && iter % gbdt_config_->bagging_freq == 0) {
Guolin Ke's avatar
Guolin Ke committed
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
    // if doesn't have query data
    if (train_data_->metadata().query_boundaries() == nullptr) {
      bag_data_cnt_ =
        static_cast<data_size_t>(gbdt_config_->bagging_fraction * num_data_);
      out_of_bag_data_cnt_ = num_data_ - bag_data_cnt_;
      data_size_t cur_left_cnt = 0;
      data_size_t cur_right_cnt = 0;
      // random bagging, minimal unit is one record
      for (data_size_t i = 0; i < num_data_; ++i) {
        double prob =
          (bag_data_cnt_ - cur_left_cnt) / static_cast<double>(num_data_ - i);
        if (random_.NextDouble() < prob) {
          bag_data_indices_[cur_left_cnt++] = i;
        } else {
          out_of_bag_data_indices_[cur_right_cnt++] = i;
        }
      }
    } else {
      // if have query data
      const data_size_t* query_boundaries = train_data_->metadata().query_boundaries();
      data_size_t num_query = train_data_->metadata().num_queries();
      data_size_t bag_query_cnt =
          static_cast<data_size_t>(num_query * gbdt_config_->bagging_fraction);
      data_size_t cur_left_query_cnt = 0;
      data_size_t cur_left_cnt = 0;
      data_size_t cur_right_cnt = 0;
      // random bagging, minimal unit is one query
      for (data_size_t i = 0; i < num_query; ++i) {
        double prob =
            (bag_query_cnt - cur_left_query_cnt) / static_cast<double>(num_query - i);
        if (random_.NextDouble() < prob) {
          for (data_size_t j = query_boundaries[i]; j < query_boundaries[i + 1]; ++j) {
            bag_data_indices_[cur_left_cnt++] = j;
          }
          cur_left_query_cnt++;
        } else {
          for (data_size_t j = query_boundaries[i]; j < query_boundaries[i + 1]; ++j) {
            out_of_bag_data_indices_[cur_right_cnt++] = j;
          }
        }
      }
      bag_data_cnt_ = cur_left_cnt;
      out_of_bag_data_cnt_ = num_data_ - bag_data_cnt_;
    }
Guolin Ke's avatar
Guolin Ke committed
204
    Log::Debug("Re-bagging, using %d data to train", bag_data_cnt_);
Guolin Ke's avatar
Guolin Ke committed
205
    // set bagging data to tree learner
Guolin Ke's avatar
Guolin Ke committed
206
    tree_learner_[curr_class]->SetBaggingData(bag_data_indices_.data(), bag_data_cnt_);
Guolin Ke's avatar
Guolin Ke committed
207
208
209
  }
}

210
void GBDT::UpdateScoreOutOfBag(const Tree* tree, const int curr_class) {
Hui Xue's avatar
Hui Xue committed
211
  // we need to predict out-of-bag socres of data for boosting
Guolin Ke's avatar
Guolin Ke committed
212
  if (!out_of_bag_data_indices_.empty()) {
Guolin Ke's avatar
Guolin Ke committed
213
    train_score_updater_->AddScore(tree, out_of_bag_data_indices_.data(), out_of_bag_data_cnt_, curr_class);
Guolin Ke's avatar
Guolin Ke committed
214
215
216
  }
}

217
bool GBDT::TrainOneIter(const score_t* gradient, const score_t* hessian, bool is_eval) {
Guolin Ke's avatar
Guolin Ke committed
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
  // boosting first
  if (gradient == nullptr || hessian == nullptr) {
    Boosting();
    gradient = gradients_.data();
    hessian = hessians_.data();
  }

  for (int curr_class = 0; curr_class < num_class_; ++curr_class) {
    // bagging logic
    Bagging(iter_, curr_class);

    // train a new tree
    std::unique_ptr<Tree> new_tree(tree_learner_[curr_class]->Train(gradient + curr_class * num_data_, hessian + curr_class * num_data_));
    // if cannot learn a new tree, then stop
    if (new_tree->num_leaves() <= 1) {
      Log::Info("Stopped training because there are no more leafs that meet the split requirements.");
      return true;
235
    }
236

Guolin Ke's avatar
Guolin Ke committed
237
238
239
240
241
    // shrinkage by learning rate
    new_tree->Shrinkage(shrinkage_rate_);
    // update score
    UpdateScore(new_tree.get(), curr_class);
    UpdateScoreOutOfBag(new_tree.get(), curr_class);
242

Guolin Ke's avatar
Guolin Ke committed
243
244
245
246
247
248
249
250
251
    // add model
    models_.push_back(std::move(new_tree));
  }
  ++iter_;
  if (is_eval) {
    return EvalAndCheckEarlyStopping();
  } else {
    return false;
  }
252

Guolin Ke's avatar
Guolin Ke committed
253
}
254

wxchan's avatar
wxchan committed
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
void GBDT::RollbackOneIter() {
  if (iter_ == 0) { return; }
  int cur_iter = iter_ + num_init_iteration_ - 1;
  // reset score
  for (int curr_class = 0; curr_class < num_class_; ++curr_class) {
    auto curr_tree = cur_iter * num_class_ + curr_class;
    models_[curr_tree]->Shrinkage(-1.0);
    train_score_updater_->AddScore(models_[curr_tree].get(), curr_class);
    for (auto& score_updater : valid_score_updater_) {
      score_updater->AddScore(models_[curr_tree].get(), curr_class);
    }
  }
  // remove model
  for (int curr_class = 0; curr_class < num_class_; ++curr_class) {
    models_.pop_back();
  }
  --iter_;
}

Guolin Ke's avatar
Guolin Ke committed
274
bool GBDT::EvalAndCheckEarlyStopping() {
275
276
  bool is_met_early_stopping = false;
  // print message for metric
Guolin Ke's avatar
Guolin Ke committed
277
278
  auto best_msg = OutputMetric(iter_);
  is_met_early_stopping = !best_msg.empty();
279
280
281
  if (is_met_early_stopping) {
    Log::Info("Early stopping at iteration %d, the best iteration round is %d",
      iter_, iter_ - early_stopping_round_);
Guolin Ke's avatar
Guolin Ke committed
282
    Log::Info("Output of best iteration round:\n%s", best_msg.c_str());
283
    // pop last early_stopping_round_ models
284
    for (int i = 0; i < early_stopping_round_ * num_class_; ++i) {
285
286
287
288
      models_.pop_back();
    }
  }
  return is_met_early_stopping;
Guolin Ke's avatar
Guolin Ke committed
289
290
}

291
void GBDT::UpdateScore(const Tree* tree, const int curr_class) {
Guolin Ke's avatar
Guolin Ke committed
292
  // update training score
Guolin Ke's avatar
Guolin Ke committed
293
  train_score_updater_->AddScore(tree_learner_[curr_class].get(), curr_class);
Guolin Ke's avatar
Guolin Ke committed
294
  // update validation score
Guolin Ke's avatar
Guolin Ke committed
295
296
  for (auto& score_updater : valid_score_updater_) {
    score_updater->AddScore(tree, curr_class);
Guolin Ke's avatar
Guolin Ke committed
297
298
299
  }
}

Guolin Ke's avatar
Guolin Ke committed
300
301
302
303
304
std::string GBDT::OutputMetric(int iter) {
  bool need_output = (iter % gbdt_config_->output_freq) == 0;
  std::string ret = "";
  std::stringstream msg_buf;
  std::vector<std::pair<int, int>> meet_early_stopping_pairs;
Guolin Ke's avatar
Guolin Ke committed
305
  // print training metric
Guolin Ke's avatar
Guolin Ke committed
306
  if (need_output) {
307
308
309
    for (auto& sub_metric : training_metrics_) {
      auto name = sub_metric->GetName();
      auto scores = sub_metric->Eval(train_score_updater_->score());
Guolin Ke's avatar
Guolin Ke committed
310
      for (size_t k = 0; k < name.size(); ++k) {
Guolin Ke's avatar
Guolin Ke committed
311
312
313
314
315
316
317
318
        std::stringstream tmp_buf;
        tmp_buf << "Iteration:" << iter
          << ", training " << name[k]
          << " : " << scores[k];
        Log::Info(tmp_buf.str().c_str());
        if (early_stopping_round_ > 0) {
          msg_buf << tmp_buf.str() << std::endl;
        }
319
      }
320
    }
Guolin Ke's avatar
Guolin Ke committed
321
322
  }
  // print validation metric
Guolin Ke's avatar
Guolin Ke committed
323
  if (need_output || early_stopping_round_ > 0) {
324
325
326
    for (size_t i = 0; i < valid_metrics_.size(); ++i) {
      for (size_t j = 0; j < valid_metrics_[i].size(); ++j) {
        auto test_scores = valid_metrics_[i][j]->Eval(valid_score_updater_[i]->score());
Guolin Ke's avatar
Guolin Ke committed
327
328
329
330
331
332
333
334
335
336
337
        auto name = valid_metrics_[i][j]->GetName();
        for (size_t k = 0; k < name.size(); ++k) {
          std::stringstream tmp_buf;
          tmp_buf << "Iteration:" << iter
            << ", valid_" << i + 1 << " " << name[k]
            << " : " << test_scores[k];
          if (need_output) {
            Log::Info(tmp_buf.str().c_str());
          }
          if (early_stopping_round_ > 0) {
            msg_buf << tmp_buf.str() << std::endl;
338
          }
wxchan's avatar
wxchan committed
339
        }
Guolin Ke's avatar
Guolin Ke committed
340
        if (ret.empty() && early_stopping_round_ > 0) {
341
342
343
          auto cur_score = valid_metrics_[i][j]->factor_to_bigger_better() * test_scores.back();
          if (cur_score > best_score_[i][j]) {
            best_score_[i][j] = cur_score;
344
            best_iter_[i][j] = iter;
Guolin Ke's avatar
Guolin Ke committed
345
            meet_early_stopping_pairs.emplace_back(i, j);
346
          } else {
Guolin Ke's avatar
Guolin Ke committed
347
            if (iter - best_iter_[i][j] >= early_stopping_round_) { ret = best_msg_[i][j]; }
348
          }
wxchan's avatar
wxchan committed
349
350
        }
      }
Guolin Ke's avatar
Guolin Ke committed
351
352
    }
  }
Guolin Ke's avatar
Guolin Ke committed
353
354
355
  for (auto& pair : meet_early_stopping_pairs) {
    best_msg_[pair.first][pair.second] = msg_buf.str();
  }
wxchan's avatar
wxchan committed
356
  return ret;
Guolin Ke's avatar
Guolin Ke committed
357
358
}

359
/*! \brief Get eval result */
360
361
362
363
std::vector<double> GBDT::GetEvalAt(int data_idx) const {
  CHECK(data_idx >= 0 && data_idx <= static_cast<int>(valid_metrics_.size()));
  std::vector<double> ret;
  if (data_idx == 0) {
364
365
    for (auto& sub_metric : training_metrics_) {
      auto scores = sub_metric->Eval(train_score_updater_->score());
366
367
368
      for (auto score : scores) {
        ret.push_back(score);
      }
369
370
    }
  }
371
372
373
374
375
376
377
  else {
    auto used_idx = data_idx - 1;
    for (size_t j = 0; j < valid_metrics_[used_idx].size(); ++j) {
      auto test_scores = valid_metrics_[used_idx][j]->Eval(valid_score_updater_[used_idx]->score());
      for (auto score : test_scores) {
        ret.push_back(score);
      }
378
379
380
381
382
    }
  }
  return ret;
}

Guolin Ke's avatar
Guolin Ke committed
383
/*! \brief Get training scores result */
384
const score_t* GBDT::GetTrainingScore(data_size_t* out_len) {
Guolin Ke's avatar
Guolin Ke committed
385
386
  *out_len = train_score_updater_->num_data() * num_class_;
  return train_score_updater_->score();
387
388
}

wxchan's avatar
wxchan committed
389
void GBDT::GetPredictAt(int data_idx, score_t* out_result, data_size_t* out_len) {
Guolin Ke's avatar
Guolin Ke committed
390
391
392
393
394
395
  CHECK(data_idx >= 0 && data_idx <= static_cast<int>(valid_metrics_.size()));
  std::vector<double> ret;

  const score_t* raw_scores = nullptr;
  data_size_t num_data = 0;
  if (data_idx == 0) {
wxchan's avatar
wxchan committed
396
    raw_scores = GetTrainingScore(out_len);
Guolin Ke's avatar
Guolin Ke committed
397
398
399
400
401
    num_data = train_score_updater_->num_data();
  } else {
    auto used_idx = data_idx - 1;
    raw_scores = valid_score_updater_[used_idx]->score();
    num_data = valid_score_updater_[used_idx]->num_data();
wxchan's avatar
wxchan committed
402
    *out_len = num_data * num_class_;
Guolin Ke's avatar
Guolin Ke committed
403
404
  }
  if (num_class_ > 1) {
wxchan's avatar
wxchan committed
405
#pragma omp parallel for schedule(static)
Guolin Ke's avatar
Guolin Ke committed
406
407
408
409
410
411
412
413
414
415
    for (data_size_t i = 0; i < num_data; ++i) {
      std::vector<double> tmp_result;
      for (int j = 0; j < num_class_; ++j) {
        tmp_result.push_back(raw_scores[j * num_data + i]);
      }
      Common::Softmax(&tmp_result);
      for (int j = 0; j < num_class_; ++j) {
        out_result[j * num_data + i] = static_cast<score_t>(tmp_result[i]);
      }
    }
Guolin Ke's avatar
Guolin Ke committed
416
  } else if(sigmoid_ > 0.0f){
wxchan's avatar
wxchan committed
417
#pragma omp parallel for schedule(static)
Guolin Ke's avatar
Guolin Ke committed
418
419
420
421
    for (data_size_t i = 0; i < num_data; ++i) {
      out_result[i] = static_cast<score_t>(1.0f / (1.0f + std::exp(-2.0f * sigmoid_ * raw_scores[i])));
    }
  } else {
wxchan's avatar
wxchan committed
422
#pragma omp parallel for schedule(static)
Guolin Ke's avatar
Guolin Ke committed
423
424
425
426
427
428
429
    for (data_size_t i = 0; i < num_data; ++i) {
      out_result[i] = raw_scores[i];
    }
  }

}

Guolin Ke's avatar
Guolin Ke committed
430
void GBDT::Boosting() {
431
432
433
  if (object_function_ == nullptr) {
    Log::Fatal("No object function provided");
  }
Hui Xue's avatar
Hui Xue committed
434
  // objective function will calculate gradients and hessians
Guolin Ke's avatar
Guolin Ke committed
435
  int num_score = 0;
Guolin Ke's avatar
Guolin Ke committed
436
  object_function_->
Guolin Ke's avatar
Guolin Ke committed
437
    GetGradients(GetTrainingScore(&num_score), gradients_.data(), hessians_.data());
Guolin Ke's avatar
Guolin Ke committed
438
439
}

wxchan's avatar
wxchan committed
440
std::string GBDT::DumpModel() const {
Guolin Ke's avatar
Guolin Ke committed
441
  std::stringstream str_buf;
wxchan's avatar
wxchan committed
442

Guolin Ke's avatar
Guolin Ke committed
443
444
445
446
447
448
  str_buf << "{";
  str_buf << "\"name\":\"" << Name() << "\"," << std::endl;
  str_buf << "\"num_class\":" << num_class_ << "," << std::endl;
  str_buf << "\"label_index\":" << label_idx_ << "," << std::endl;
  str_buf << "\"max_feature_idx\":" << max_feature_idx_ << "," << std::endl;
  str_buf << "\"sigmoid\":" << sigmoid_ << "," << std::endl;
wxchan's avatar
wxchan committed
449

Guolin Ke's avatar
Guolin Ke committed
450
451
452
453
454
455
  // output feature names
  auto feature_names = std::ref(feature_names_);
  if (train_data_ != nullptr) {
    feature_names = std::ref(train_data_->feature_names());
  }

Guolin Ke's avatar
Guolin Ke committed
456
  str_buf << "\"feature_names\":[\"" 
Guolin Ke's avatar
Guolin Ke committed
457
458
459
     << Common::Join(feature_names.get(), "\",\"") << "\"]," 
     << std::endl;

Guolin Ke's avatar
Guolin Ke committed
460
  str_buf << "\"tree_info\":[";
wxchan's avatar
wxchan committed
461
462
  for (int i = 0; i < static_cast<int>(models_.size()); ++i) {
    if (i > 0) {
Guolin Ke's avatar
Guolin Ke committed
463
      str_buf << ",";
wxchan's avatar
wxchan committed
464
    }
Guolin Ke's avatar
Guolin Ke committed
465
466
467
468
    str_buf << "{";
    str_buf << "\"tree_index\":" << i << ",";
    str_buf << models_[i]->ToJSON();
    str_buf << "}";
wxchan's avatar
wxchan committed
469
  }
Guolin Ke's avatar
Guolin Ke committed
470
  str_buf << "]" << std::endl;
wxchan's avatar
wxchan committed
471

Guolin Ke's avatar
Guolin Ke committed
472
  str_buf << "}" << std::endl;
wxchan's avatar
wxchan committed
473

Guolin Ke's avatar
Guolin Ke committed
474
  return str_buf.str();
wxchan's avatar
wxchan committed
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
}

void GBDT::SaveModelToFile(int num_iteration, const char* filename) const {
  /*! \brief File to write models */
  std::ofstream output_file;
  output_file.open(filename);
  // output model type
  output_file << Name() << std::endl;
  // output number of class
  output_file << "num_class=" << num_class_ << std::endl;
  // output label index
  output_file << "label_index=" << label_idx_ << std::endl;
  // output max_feature_idx
  output_file << "max_feature_idx=" << max_feature_idx_ << std::endl;
  // output objective name
  if (object_function_ != nullptr) {
    output_file << "objective=" << object_function_->GetName() << std::endl;
492
  }
wxchan's avatar
wxchan committed
493
494
  // output sigmoid parameter
  output_file << "sigmoid=" << sigmoid_ << std::endl;
Guolin Ke's avatar
Guolin Ke committed
495
496
497
498
499
500
  // output feature names
  auto feature_names = std::ref(feature_names_);
  if (train_data_ != nullptr) {
    feature_names = std::ref(train_data_->feature_names());
  }
  output_file << "feature_names=" << Common::Join(feature_names.get(), " ") << std::endl;
wxchan's avatar
wxchan committed
501

Guolin Ke's avatar
Guolin Ke committed
502
  output_file << std::endl;
wxchan's avatar
wxchan committed
503
504
  int num_used_model = 0;
  if (num_iteration <= 0) {
Guolin Ke's avatar
Guolin Ke committed
505
506
    num_used_model = static_cast<int>(models_.size());
  } else {
wxchan's avatar
wxchan committed
507
    num_used_model = num_iteration * num_class_;
Guolin Ke's avatar
Guolin Ke committed
508
  }
wxchan's avatar
wxchan committed
509
  num_used_model = std::min(num_used_model, static_cast<int>(models_.size()));
510
  // output tree models
wxchan's avatar
wxchan committed
511
512
513
  for (int i = 0; i < num_used_model; ++i) {
    output_file << "Tree=" << i << std::endl;
    output_file << models_[i]->ToString() << std::endl;
514
  }
515

wxchan's avatar
wxchan committed
516
517
518
519
  std::vector<std::pair<size_t, std::string>> pairs = FeatureImportance();
  output_file << std::endl << "feature importances:" << std::endl;
  for (size_t i = 0; i < pairs.size(); ++i) {
    output_file << pairs[i].second << "=" << std::to_string(pairs[i].first) << std::endl;
Guolin Ke's avatar
Guolin Ke committed
520
  }
wxchan's avatar
wxchan committed
521
  output_file.close();
Guolin Ke's avatar
Guolin Ke committed
522
523
}

Guolin Ke's avatar
Guolin Ke committed
524
void GBDT::LoadModelFromString(const std::string& model_str) {
Guolin Ke's avatar
Guolin Ke committed
525
526
527
  // use serialized string to restore this object
  models_.clear();
  std::vector<std::string> lines = Common::Split(model_str.c_str(), '\n');
528
529

  // get number of classes
530
531
532
533
  auto line = Common::FindFromLines(lines, "num_class=");
  if (line.size() > 0) {
    Common::Atoi(Common::Split(line.c_str(), '=')[1].c_str(), &num_class_);
  } else {
534
    Log::Fatal("Model file doesn't specify the number of classes");
535
536
    return;
  }
Guolin Ke's avatar
Guolin Ke committed
537
  // get index of label
538
539
540
541
  line = Common::FindFromLines(lines, "label_index=");
  if (line.size() > 0) {
    Common::Atoi(Common::Split(line.c_str(), '=')[1].c_str(), &label_idx_);
  } else {
542
    Log::Fatal("Model file doesn't specify the label index");
Guolin Ke's avatar
Guolin Ke committed
543
544
    return;
  }
Guolin Ke's avatar
Guolin Ke committed
545
  // get max_feature_idx first
546
547
548
549
  line = Common::FindFromLines(lines, "max_feature_idx=");
  if (line.size() > 0) {
    Common::Atoi(Common::Split(line.c_str(), '=')[1].c_str(), &max_feature_idx_);
  } else {
550
    Log::Fatal("Model file doesn't specify max_feature_idx");
Guolin Ke's avatar
Guolin Ke committed
551
552
553
    return;
  }
  // get sigmoid parameter
554
555
556
557
  line = Common::FindFromLines(lines, "sigmoid=");
  if (line.size() > 0) {
    Common::Atof(Common::Split(line.c_str(), '=')[1].c_str(), &sigmoid_);
  } else {
558
    sigmoid_ = -1.0f;
Guolin Ke's avatar
Guolin Ke committed
559
  }
Guolin Ke's avatar
Guolin Ke committed
560
561
562
563
564
565
566
567
568
569
570
571
572
  // get feature names
  line = Common::FindFromLines(lines, "feature_names=");
  if (line.size() > 0) {
    feature_names_ = Common::Split(Common::Split(line.c_str(), '=')[1].c_str(), ' ');
    if (feature_names_.size() != static_cast<size_t>(max_feature_idx_ + 1)) {
      Log::Fatal("Wrong size of feature_names");
      return;
    }
  } else {
    Log::Fatal("Model file doesn't contain feature names");
    return;
  }

Guolin Ke's avatar
Guolin Ke committed
573
  // get tree models
574
  size_t i = 0;
Guolin Ke's avatar
Guolin Ke committed
575
576
577
578
579
580
581
  while (i < lines.size()) {
    size_t find_pos = lines[i].find("Tree=");
    if (find_pos != std::string::npos) {
      ++i;
      int start = static_cast<int>(i);
      while (i < lines.size() && lines[i].find("Tree=") == std::string::npos) { ++i; }
      int end = static_cast<int>(i);
Guolin Ke's avatar
Guolin Ke committed
582
      std::string tree_str = Common::Join<std::string>(lines, start, end, "\n");
Guolin Ke's avatar
Guolin Ke committed
583
584
      auto new_tree = std::unique_ptr<Tree>(new Tree(tree_str));
      models_.push_back(std::move(new_tree));
Guolin Ke's avatar
Guolin Ke committed
585
586
587
588
    } else {
      ++i;
    }
  }
589
  Log::Info("Finished loading %d models", models_.size());
wxchan's avatar
wxchan committed
590
591
  num_iteration_for_pred_ = static_cast<int>(models_.size()) / num_class_;
  num_init_iteration_ = num_iteration_for_pred_;
Guolin Ke's avatar
Guolin Ke committed
592
593
}

wxchan's avatar
wxchan committed
594
std::vector<std::pair<size_t, std::string>> GBDT::FeatureImportance() const {
Guolin Ke's avatar
Guolin Ke committed
595
596
597
598
  auto feature_names = std::ref(feature_names_);
  if (train_data_ != nullptr) {
    feature_names = std::ref(train_data_->feature_names());
  }
599
  std::vector<size_t> feature_importances(max_feature_idx_ + 1, 0);
600
    for (size_t iter = 0; iter < models_.size(); ++iter) {
601
602
        for (int split_idx = 0; split_idx < models_[iter]->num_leaves() - 1; ++split_idx) {
            ++feature_importances[models_[iter]->split_feature_real(split_idx)];
wxchan's avatar
wxchan committed
603
604
        }
    }
605
606
607
    // store the importance first
    std::vector<std::pair<size_t, std::string>> pairs;
    for (size_t i = 0; i < feature_importances.size(); ++i) {
Guolin Ke's avatar
Guolin Ke committed
608
      if (feature_importances[i] > 0) {
Guolin Ke's avatar
Guolin Ke committed
609
        pairs.emplace_back(feature_importances[i], feature_names.get().at(i));
Guolin Ke's avatar
Guolin Ke committed
610
      }
611
612
613
614
615
    }
    // sort the importance
    std::sort(pairs.begin(), pairs.end(),
      [](const std::pair<size_t, std::string>& lhs,
        const std::pair<size_t, std::string>& rhs) {
616
      return lhs.first > rhs.first;
617
    });
wxchan's avatar
wxchan committed
618
    return pairs;
wxchan's avatar
wxchan committed
619
620
}

621
622
std::vector<double> GBDT::PredictRaw(const double* value) const {
  std::vector<double> ret(num_class_, 0.0f);
wxchan's avatar
wxchan committed
623
  for (int i = 0; i < num_iteration_for_pred_; ++i) {
624
625
626
    for (int j = 0; j < num_class_; ++j) {
      ret[j] += models_[i * num_class_ + j]->Predict(value);
    }
Guolin Ke's avatar
Guolin Ke committed
627
628
629
630
  }
  return ret;
}

631
std::vector<double> GBDT::Predict(const double* value) const {
632
  std::vector<double> ret(num_class_, 0.0f);
wxchan's avatar
wxchan committed
633
  for (int i = 0; i < num_iteration_for_pred_; ++i) {
634
635
    for (int j = 0; j < num_class_; ++j) {
      ret[j] += models_[i * num_class_ + j]->Predict(value);
636
637
    }
  }
638
639
640
641
642
643
  // if need sigmoid transform
  if (sigmoid_ > 0 && num_class_ == 1) {
    ret[0] = 1.0f / (1.0f + std::exp(- 2.0f * sigmoid_ * ret[0]));
  } else if (num_class_ > 1) {
    Common::Softmax(&ret);
  }
644
645
646
  return ret;
}

647
std::vector<int> GBDT::PredictLeafIndex(const double* value) const {
wxchan's avatar
wxchan committed
648
  std::vector<int> ret;
wxchan's avatar
wxchan committed
649
  for (int i = 0; i < num_iteration_for_pred_; ++i) {
650
651
652
    for (int j = 0; j < num_class_; ++j) {
      ret.push_back(models_[i * num_class_ + j]->PredictLeafIndex(value));
    }
wxchan's avatar
wxchan committed
653
654
655
656
  }
  return ret;
}

Guolin Ke's avatar
Guolin Ke committed
657
}  // namespace LightGBM