gbdt.cpp 22.3 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
35
  num_class_ = config->num_class;
  train_data_ = nullptr;
Guolin Ke's avatar
Guolin Ke committed
36
  gbdt_config_ = nullptr;
wxchan's avatar
wxchan committed
37
38
39
40
41
  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
42
43
  if (train_data == nullptr) { return; }
  auto new_config = std::unique_ptr<BoostingConfig>(new BoostingConfig(*config));
wxchan's avatar
wxchan committed
44
45
46
  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
47
48
  early_stopping_round_ = new_config->early_stopping_round;
  shrinkage_rate_ = new_config->learning_rate;
49
50
51
52
  // cannot reset seed, only create one time
  if (gbdt_config_ == nullptr ) {
    random_ = Random(new_config->bagging_seed);
  }
Guolin Ke's avatar
Guolin Ke committed
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69

  // create tree learner, only create once
  if (gbdt_config_ == nullptr) {
    tree_learner_.clear();
    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
  if (train_data_ != train_data) {
    for (int i = 0; i < num_class_; ++i) {
      tree_learner_[i]->Init(train_data);
    }
  }
  // reset config for tree learner
Guolin Ke's avatar
Guolin Ke committed
70
  for (int i = 0; i < num_class_; ++i) {
Guolin Ke's avatar
Guolin Ke committed
71
    tree_learner_[i]->ResetConfig(&new_config->tree_config);
72
  }
Guolin Ke's avatar
Guolin Ke committed
73

Guolin Ke's avatar
Guolin Ke committed
74
  object_function_ = object_function;
Guolin Ke's avatar
Guolin Ke committed
75

Guolin Ke's avatar
Guolin Ke committed
76
  sigmoid_ = -1.0f;
wxchan's avatar
wxchan committed
77
  if (object_function_ != nullptr
Guolin Ke's avatar
Guolin Ke committed
78
79
    && std::string(object_function_->GetName()) == std::string("binary")) {
    // only binary classification need sigmoid transform
Guolin Ke's avatar
Guolin Ke committed
80
    sigmoid_ = new_config->sigmoid;
81
  }
Guolin Ke's avatar
Guolin Ke committed
82

wxchan's avatar
wxchan committed
83
  if (train_data_ != train_data) {
Guolin Ke's avatar
Guolin Ke committed
84
85
86
87
88
89
    // 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
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
    // 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) {
      gradients_ = std::vector<score_t>(num_data_ * num_class_);
      hessians_ = std::vector<score_t>(num_data_ * num_class_);
    }
    // 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
110
111
112
113
114
  }

  if (train_data_ != train_data
      || gbdt_config_ == nullptr
      || (gbdt_config_->bagging_fraction != new_config->bagging_fraction)) {
wxchan's avatar
wxchan committed
115
    // if need bagging, create buffer
Guolin Ke's avatar
Guolin Ke committed
116
    if (new_config->bagging_fraction < 1.0 && new_config->bagging_freq > 0) {
wxchan's avatar
wxchan committed
117
118
119
120
121
122
123
124
      out_of_bag_data_indices_ = std::vector<data_size_t>(num_data_);
      bag_data_indices_ = std::vector<data_size_t>(num_data_);
    } 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
125
  }
wxchan's avatar
wxchan committed
126
  train_data_ = train_data;
Guolin Ke's avatar
Guolin Ke committed
127
  gbdt_config_.reset(new_config.release());
Guolin Ke's avatar
Guolin Ke committed
128
129
}

wxchan's avatar
wxchan committed
130
void GBDT::AddValidDataset(const Dataset* valid_data,
Guolin Ke's avatar
Guolin Ke committed
131
  const std::vector<const Metric*>& valid_metrics) {
wxchan's avatar
wxchan committed
132
133
  if (!train_data_->CheckAlign(*valid_data)) {
    Log::Fatal("cannot add validation data, since it has different bin mappers with training data");
134
  }
Guolin Ke's avatar
Guolin Ke committed
135
  // for a validation dataset, we need its score and metric
Guolin Ke's avatar
Guolin Ke committed
136
  auto new_score_updater = std::unique_ptr<ScoreUpdater>(new ScoreUpdater(valid_data, num_class_));
wxchan's avatar
wxchan committed
137
138
139
140
141
142
143
  // 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
144
  valid_score_updater_.push_back(std::move(new_score_updater));
Guolin Ke's avatar
Guolin Ke committed
145
  valid_metrics_.emplace_back();
146
147
148
  if (early_stopping_round_ > 0) {
    best_iter_.emplace_back();
    best_score_.emplace_back();
Guolin Ke's avatar
Guolin Ke committed
149
    best_msg_.emplace_back();
150
  }
Guolin Ke's avatar
Guolin Ke committed
151
152
  for (const auto& metric : valid_metrics) {
    valid_metrics_.back().push_back(metric);
153
154
155
    if (early_stopping_round_ > 0) {
      best_iter_.back().push_back(0);
      best_score_.back().push_back(kMinScore);
Guolin Ke's avatar
Guolin Ke committed
156
      best_msg_.back().emplace_back();
157
    }
Guolin Ke's avatar
Guolin Ke committed
158
  }
Guolin Ke's avatar
Guolin Ke committed
159
  valid_metrics_.back().shrink_to_fit();
Guolin Ke's avatar
Guolin Ke committed
160
161
162
}


163
void GBDT::Bagging(int iter, const int curr_class) {
Guolin Ke's avatar
Guolin Ke committed
164
  // if need bagging
Guolin Ke's avatar
Guolin Ke committed
165
  if (!out_of_bag_data_indices_.empty() && iter % gbdt_config_->bagging_freq == 0) {
Guolin Ke's avatar
Guolin Ke committed
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
204
205
206
207
208
209
    // 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
210
    Log::Debug("Re-bagging, using %d data to train", bag_data_cnt_);
Guolin Ke's avatar
Guolin Ke committed
211
    // set bagging data to tree learner
Guolin Ke's avatar
Guolin Ke committed
212
    tree_learner_[curr_class]->SetBaggingData(bag_data_indices_.data(), bag_data_cnt_);
Guolin Ke's avatar
Guolin Ke committed
213
214
215
  }
}

216
void GBDT::UpdateScoreOutOfBag(const Tree* tree, const int curr_class) {
Hui Xue's avatar
Hui Xue committed
217
  // we need to predict out-of-bag socres of data for boosting
Guolin Ke's avatar
Guolin Ke committed
218
  if (!out_of_bag_data_indices_.empty()) {
Guolin Ke's avatar
Guolin Ke committed
219
    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
220
221
222
  }
}

223
bool GBDT::TrainOneIter(const score_t* gradient, const score_t* hessian, bool is_eval) {
Guolin Ke's avatar
Guolin Ke committed
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
  // 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;
241
    }
242

Guolin Ke's avatar
Guolin Ke committed
243
244
245
246
247
    // shrinkage by learning rate
    new_tree->Shrinkage(shrinkage_rate_);
    // update score
    UpdateScore(new_tree.get(), curr_class);
    UpdateScoreOutOfBag(new_tree.get(), curr_class);
248

Guolin Ke's avatar
Guolin Ke committed
249
250
251
252
253
254
255
256
257
    // add model
    models_.push_back(std::move(new_tree));
  }
  ++iter_;
  if (is_eval) {
    return EvalAndCheckEarlyStopping();
  } else {
    return false;
  }
258

Guolin Ke's avatar
Guolin Ke committed
259
}
260

wxchan's avatar
wxchan committed
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
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
280
bool GBDT::EvalAndCheckEarlyStopping() {
281
282
  bool is_met_early_stopping = false;
  // print message for metric
Guolin Ke's avatar
Guolin Ke committed
283
284
  auto best_msg = OutputMetric(iter_);
  is_met_early_stopping = !best_msg.empty();
285
286
287
  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
288
    Log::Info("Output of best iteration round:\n%s", best_msg.c_str());
289
    // pop last early_stopping_round_ models
290
    for (int i = 0; i < early_stopping_round_ * num_class_; ++i) {
291
292
293
294
      models_.pop_back();
    }
  }
  return is_met_early_stopping;
Guolin Ke's avatar
Guolin Ke committed
295
296
}

297
void GBDT::UpdateScore(const Tree* tree, const int curr_class) {
Guolin Ke's avatar
Guolin Ke committed
298
  // update training score
Guolin Ke's avatar
Guolin Ke committed
299
  train_score_updater_->AddScore(tree_learner_[curr_class].get(), curr_class);
Guolin Ke's avatar
Guolin Ke committed
300
  // update validation score
Guolin Ke's avatar
Guolin Ke committed
301
302
  for (auto& score_updater : valid_score_updater_) {
    score_updater->AddScore(tree, curr_class);
Guolin Ke's avatar
Guolin Ke committed
303
304
305
  }
}

Guolin Ke's avatar
Guolin Ke committed
306
307
308
309
310
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
311
  // print training metric
Guolin Ke's avatar
Guolin Ke committed
312
  if (need_output) {
313
314
315
    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
316
      for (size_t k = 0; k < name.size(); ++k) {
Guolin Ke's avatar
Guolin Ke committed
317
318
319
320
321
322
323
324
        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;
        }
325
      }
326
    }
Guolin Ke's avatar
Guolin Ke committed
327
328
  }
  // print validation metric
Guolin Ke's avatar
Guolin Ke committed
329
  if (need_output || early_stopping_round_ > 0) {
330
331
332
    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
333
334
335
336
337
338
339
340
341
342
343
        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;
344
          }
wxchan's avatar
wxchan committed
345
        }
Guolin Ke's avatar
Guolin Ke committed
346
        if (ret.empty() && early_stopping_round_ > 0) {
347
348
349
          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;
350
            best_iter_[i][j] = iter;
Guolin Ke's avatar
Guolin Ke committed
351
            meet_early_stopping_pairs.emplace_back(i, j);
352
          } else {
Guolin Ke's avatar
Guolin Ke committed
353
            if (iter - best_iter_[i][j] >= early_stopping_round_) { ret = best_msg_[i][j]; }
354
          }
wxchan's avatar
wxchan committed
355
356
        }
      }
Guolin Ke's avatar
Guolin Ke committed
357
358
    }
  }
Guolin Ke's avatar
Guolin Ke committed
359
360
361
  for (auto& pair : meet_early_stopping_pairs) {
    best_msg_[pair.first][pair.second] = msg_buf.str();
  }
wxchan's avatar
wxchan committed
362
  return ret;
Guolin Ke's avatar
Guolin Ke committed
363
364
}

365
/*! \brief Get eval result */
366
367
368
369
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) {
370
371
    for (auto& sub_metric : training_metrics_) {
      auto scores = sub_metric->Eval(train_score_updater_->score());
372
373
374
      for (auto score : scores) {
        ret.push_back(score);
      }
375
376
    }
  }
377
378
379
380
381
382
383
  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);
      }
384
385
386
387
388
    }
  }
  return ret;
}

Guolin Ke's avatar
Guolin Ke committed
389
/*! \brief Get training scores result */
390
const score_t* GBDT::GetTrainingScore(data_size_t* out_len) {
Guolin Ke's avatar
Guolin Ke committed
391
392
  *out_len = train_score_updater_->num_data() * num_class_;
  return train_score_updater_->score();
393
394
}

wxchan's avatar
wxchan committed
395
void GBDT::GetPredictAt(int data_idx, score_t* out_result, data_size_t* out_len) {
Guolin Ke's avatar
Guolin Ke committed
396
397
398
399
400
401
  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
402
    raw_scores = GetTrainingScore(out_len);
Guolin Ke's avatar
Guolin Ke committed
403
404
405
406
407
    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
408
    *out_len = num_data * num_class_;
Guolin Ke's avatar
Guolin Ke committed
409
410
  }
  if (num_class_ > 1) {
wxchan's avatar
wxchan committed
411
#pragma omp parallel for schedule(static)
Guolin Ke's avatar
Guolin Ke committed
412
413
414
415
416
417
418
419
420
421
    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
422
  } else if(sigmoid_ > 0.0f){
wxchan's avatar
wxchan committed
423
#pragma omp parallel for schedule(static)
Guolin Ke's avatar
Guolin Ke committed
424
425
426
427
    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
428
#pragma omp parallel for schedule(static)
Guolin Ke's avatar
Guolin Ke committed
429
430
431
432
433
434
435
    for (data_size_t i = 0; i < num_data; ++i) {
      out_result[i] = raw_scores[i];
    }
  }

}

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

wxchan's avatar
wxchan committed
446
std::string GBDT::DumpModel() const {
Guolin Ke's avatar
Guolin Ke committed
447
  std::stringstream str_buf;
wxchan's avatar
wxchan committed
448

Guolin Ke's avatar
Guolin Ke committed
449
450
451
452
453
454
  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
455

Guolin Ke's avatar
Guolin Ke committed
456
457
458
459
460
461
  // 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
462
  str_buf << "\"feature_names\":[\"" 
Guolin Ke's avatar
Guolin Ke committed
463
464
465
     << Common::Join(feature_names.get(), "\",\"") << "\"]," 
     << std::endl;

Guolin Ke's avatar
Guolin Ke committed
466
  str_buf << "\"tree_info\":[";
wxchan's avatar
wxchan committed
467
468
  for (int i = 0; i < static_cast<int>(models_.size()); ++i) {
    if (i > 0) {
Guolin Ke's avatar
Guolin Ke committed
469
      str_buf << ",";
wxchan's avatar
wxchan committed
470
    }
Guolin Ke's avatar
Guolin Ke committed
471
472
473
474
    str_buf << "{";
    str_buf << "\"tree_index\":" << i << ",";
    str_buf << models_[i]->ToJSON();
    str_buf << "}";
wxchan's avatar
wxchan committed
475
  }
Guolin Ke's avatar
Guolin Ke committed
476
  str_buf << "]" << std::endl;
wxchan's avatar
wxchan committed
477

Guolin Ke's avatar
Guolin Ke committed
478
  str_buf << "}" << std::endl;
wxchan's avatar
wxchan committed
479

Guolin Ke's avatar
Guolin Ke committed
480
  return str_buf.str();
wxchan's avatar
wxchan committed
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
}

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;
498
  }
wxchan's avatar
wxchan committed
499
500
  // output sigmoid parameter
  output_file << "sigmoid=" << sigmoid_ << std::endl;
Guolin Ke's avatar
Guolin Ke committed
501
502
503
504
505
506
  // 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
507

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

wxchan's avatar
wxchan committed
522
523
524
525
  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
526
  }
wxchan's avatar
wxchan committed
527
  output_file.close();
Guolin Ke's avatar
Guolin Ke committed
528
529
}

Guolin Ke's avatar
Guolin Ke committed
530
void GBDT::LoadModelFromString(const std::string& model_str) {
Guolin Ke's avatar
Guolin Ke committed
531
532
533
  // use serialized string to restore this object
  models_.clear();
  std::vector<std::string> lines = Common::Split(model_str.c_str(), '\n');
534
535

  // get number of classes
536
537
538
539
  auto line = Common::FindFromLines(lines, "num_class=");
  if (line.size() > 0) {
    Common::Atoi(Common::Split(line.c_str(), '=')[1].c_str(), &num_class_);
  } else {
540
    Log::Fatal("Model file doesn't specify the number of classes");
541
542
    return;
  }
Guolin Ke's avatar
Guolin Ke committed
543
  // get index of label
544
545
546
547
  line = Common::FindFromLines(lines, "label_index=");
  if (line.size() > 0) {
    Common::Atoi(Common::Split(line.c_str(), '=')[1].c_str(), &label_idx_);
  } else {
548
    Log::Fatal("Model file doesn't specify the label index");
Guolin Ke's avatar
Guolin Ke committed
549
550
    return;
  }
Guolin Ke's avatar
Guolin Ke committed
551
  // get max_feature_idx first
552
553
554
555
  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 {
556
    Log::Fatal("Model file doesn't specify max_feature_idx");
Guolin Ke's avatar
Guolin Ke committed
557
558
559
    return;
  }
  // get sigmoid parameter
560
561
562
563
  line = Common::FindFromLines(lines, "sigmoid=");
  if (line.size() > 0) {
    Common::Atof(Common::Split(line.c_str(), '=')[1].c_str(), &sigmoid_);
  } else {
564
    sigmoid_ = -1.0f;
Guolin Ke's avatar
Guolin Ke committed
565
  }
Guolin Ke's avatar
Guolin Ke committed
566
567
568
569
570
571
572
573
574
575
576
577
578
  // 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
579
  // get tree models
580
  size_t i = 0;
Guolin Ke's avatar
Guolin Ke committed
581
582
583
584
585
586
587
  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
588
      std::string tree_str = Common::Join<std::string>(lines, start, end, "\n");
Guolin Ke's avatar
Guolin Ke committed
589
590
      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
591
592
593
594
    } else {
      ++i;
    }
  }
595
  Log::Info("Finished loading %d models", models_.size());
wxchan's avatar
wxchan committed
596
597
  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
598
599
}

wxchan's avatar
wxchan committed
600
std::vector<std::pair<size_t, std::string>> GBDT::FeatureImportance() const {
Guolin Ke's avatar
Guolin Ke committed
601
602
603
604
  auto feature_names = std::ref(feature_names_);
  if (train_data_ != nullptr) {
    feature_names = std::ref(train_data_->feature_names());
  }
605
  std::vector<size_t> feature_importances(max_feature_idx_ + 1, 0);
606
    for (size_t iter = 0; iter < models_.size(); ++iter) {
607
608
        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
609
610
        }
    }
611
612
613
    // 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
614
      if (feature_importances[i] > 0) {
Guolin Ke's avatar
Guolin Ke committed
615
        pairs.emplace_back(feature_importances[i], feature_names.get().at(i));
Guolin Ke's avatar
Guolin Ke committed
616
      }
617
618
619
620
621
    }
    // 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) {
622
      return lhs.first > rhs.first;
623
    });
wxchan's avatar
wxchan committed
624
    return pairs;
wxchan's avatar
wxchan committed
625
626
}

627
628
std::vector<double> GBDT::PredictRaw(const double* value) const {
  std::vector<double> ret(num_class_, 0.0f);
wxchan's avatar
wxchan committed
629
  for (int i = 0; i < num_iteration_for_pred_; ++i) {
630
631
632
    for (int j = 0; j < num_class_; ++j) {
      ret[j] += models_[i * num_class_ + j]->Predict(value);
    }
Guolin Ke's avatar
Guolin Ke committed
633
634
635
636
  }
  return ret;
}

637
std::vector<double> GBDT::Predict(const double* value) const {
638
  std::vector<double> ret(num_class_, 0.0f);
wxchan's avatar
wxchan committed
639
  for (int i = 0; i < num_iteration_for_pred_; ++i) {
640
641
    for (int j = 0; j < num_class_; ++j) {
      ret[j] += models_[i * num_class_ + j]->Predict(value);
642
643
    }
  }
644
645
646
647
648
649
  // 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);
  }
650
651
652
  return ret;
}

653
std::vector<int> GBDT::PredictLeafIndex(const double* value) const {
wxchan's avatar
wxchan committed
654
  std::vector<int> ret;
wxchan's avatar
wxchan committed
655
  for (int i = 0; i < num_iteration_for_pred_; ++i) {
656
657
658
    for (int j = 0; j < num_class_; ++j) {
      ret.push_back(models_[i * num_class_ + j]->PredictLeafIndex(value));
    }
wxchan's avatar
wxchan committed
659
660
661
662
  }
  return ret;
}

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