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
  CHECK(data_idx >= 0 && data_idx <= static_cast<int>(valid_metrics_.size()));

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

}

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

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

Guolin Ke's avatar
Guolin Ke committed
442
443
444
445
446
447
  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
448

Guolin Ke's avatar
Guolin Ke committed
449
450
451
452
453
454
  // 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
455
  str_buf << "\"feature_names\":[\"" 
Guolin Ke's avatar
Guolin Ke committed
456
457
458
     << Common::Join(feature_names.get(), "\",\"") << "\"]," 
     << std::endl;

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

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

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

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;
491
  }
wxchan's avatar
wxchan committed
492
493
  // output sigmoid parameter
  output_file << "sigmoid=" << sigmoid_ << std::endl;
Guolin Ke's avatar
Guolin Ke committed
494
495
496
497
498
499
  // 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
500

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

wxchan's avatar
wxchan committed
515
516
517
518
  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
519
  }
wxchan's avatar
wxchan committed
520
  output_file.close();
Guolin Ke's avatar
Guolin Ke committed
521
522
}

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

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

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

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

630
std::vector<double> GBDT::Predict(const double* value) const {
631
  std::vector<double> ret(num_class_, 0.0f);
wxchan's avatar
wxchan committed
632
  for (int i = 0; i < num_iteration_for_pred_; ++i) {
633
634
    for (int j = 0; j < num_class_; ++j) {
      ret[j] += models_[i * num_class_ + j]->Predict(value);
635
636
    }
  }
637
638
639
640
641
642
  // 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);
  }
643
644
645
  return ret;
}

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

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