gbdt.cpp 21.9 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;
38
  tree_learner_ = nullptr;
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
  if (train_data_ != train_data && train_data != nullptr) {
61
62
    if (tree_learner_ == nullptr) {
      tree_learner_ = std::unique_ptr<TreeLearner>(TreeLearner::CreateTreeLearner(new_config->tree_learner_type, &new_config->tree_config));
Guolin Ke's avatar
Guolin Ke committed
63
64
    }
    // init tree learner
65
    tree_learner_->Init(train_data);
Guolin Ke's avatar
Guolin Ke committed
66

Guolin Ke's avatar
Guolin Ke committed
67
68
69
70
71
72
    // 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
73
74
75
76
77
78
79
80
81
82
83
84
85
    // 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) {
86
87
88
      size_t total_size = static_cast<size_t>(num_data_) * num_class_;
      gradients_.resize(total_size);
      hessians_.resize(total_size);
wxchan's avatar
wxchan committed
89
90
91
92
93
    }
    // 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
94
95
  }

Guolin Ke's avatar
Guolin Ke committed
96
97
  if ((train_data_ != train_data && train_data != nullptr)
    || (gbdt_config_ != nullptr && gbdt_config_->bagging_fraction != new_config->bagging_fraction)) {
wxchan's avatar
wxchan committed
98
    // if need bagging, create buffer
Guolin Ke's avatar
Guolin Ke committed
99
    if (new_config->bagging_fraction < 1.0 && new_config->bagging_freq > 0) {
100
101
      out_of_bag_data_indices_.resize(num_data_);
      bag_data_indices_.resize(num_data_);
wxchan's avatar
wxchan committed
102
103
104
105
106
107
    } 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
108
  }
wxchan's avatar
wxchan committed
109
  train_data_ = train_data;
Guolin Ke's avatar
Guolin Ke committed
110
111
  if (train_data_ != nullptr) {
    // reset config for tree learner
112
    tree_learner_->ResetConfig(&new_config->tree_config);
Guolin Ke's avatar
Guolin Ke committed
113
  }
Guolin Ke's avatar
Guolin Ke committed
114
  gbdt_config_.reset(new_config.release());
Guolin Ke's avatar
Guolin Ke committed
115
116
}

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


150
void GBDT::Bagging(int iter) {
Guolin Ke's avatar
Guolin Ke committed
151
  // if need bagging
Guolin Ke's avatar
Guolin Ke committed
152
  if (!out_of_bag_data_indices_.empty() && iter % gbdt_config_->bagging_freq == 0) {
Guolin Ke's avatar
Guolin Ke committed
153
154
155
156
157
158
159
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
    // 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
197
    Log::Debug("Re-bagging, using %d data to train", bag_data_cnt_);
Guolin Ke's avatar
Guolin Ke committed
198
    // set bagging data to tree learner
199
    tree_learner_->SetBaggingData(bag_data_indices_.data(), bag_data_cnt_);
Guolin Ke's avatar
Guolin Ke committed
200
201
202
  }
}

203
void GBDT::UpdateScoreOutOfBag(const Tree* tree, const int curr_class) {
Hui Xue's avatar
Hui Xue committed
204
  // we need to predict out-of-bag socres of data for boosting
Guolin Ke's avatar
Guolin Ke committed
205
  if (!out_of_bag_data_indices_.empty()) {
Guolin Ke's avatar
Guolin Ke committed
206
    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
207
208
209
  }
}

210
bool GBDT::TrainOneIter(const score_t* gradient, const score_t* hessian, bool is_eval) {
Guolin Ke's avatar
Guolin Ke committed
211
212
213
214
215
216
  // boosting first
  if (gradient == nullptr || hessian == nullptr) {
    Boosting();
    gradient = gradients_.data();
    hessian = hessians_.data();
  }
217
218
  // bagging logic
  Bagging(iter_);
Guolin Ke's avatar
Guolin Ke committed
219
220
221
  for (int curr_class = 0; curr_class < num_class_; ++curr_class) {

    // train a new tree
222
    std::unique_ptr<Tree> new_tree(tree_learner_->Train(gradient + curr_class * num_data_, hessian + curr_class * num_data_));
Guolin Ke's avatar
Guolin Ke committed
223
224
225
226
    // 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;
227
    }
228

Guolin Ke's avatar
Guolin Ke committed
229
230
231
232
233
    // shrinkage by learning rate
    new_tree->Shrinkage(shrinkage_rate_);
    // update score
    UpdateScore(new_tree.get(), curr_class);
    UpdateScoreOutOfBag(new_tree.get(), curr_class);
234

Guolin Ke's avatar
Guolin Ke committed
235
236
237
238
239
240
241
242
243
    // add model
    models_.push_back(std::move(new_tree));
  }
  ++iter_;
  if (is_eval) {
    return EvalAndCheckEarlyStopping();
  } else {
    return false;
  }
244

Guolin Ke's avatar
Guolin Ke committed
245
}
246

wxchan's avatar
wxchan committed
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
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
266
bool GBDT::EvalAndCheckEarlyStopping() {
267
268
  bool is_met_early_stopping = false;
  // print message for metric
Guolin Ke's avatar
Guolin Ke committed
269
270
  auto best_msg = OutputMetric(iter_);
  is_met_early_stopping = !best_msg.empty();
271
272
273
  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
274
    Log::Info("Output of best iteration round:\n%s", best_msg.c_str());
275
    // pop last early_stopping_round_ models
276
    for (int i = 0; i < early_stopping_round_ * num_class_; ++i) {
277
278
279
280
      models_.pop_back();
    }
  }
  return is_met_early_stopping;
Guolin Ke's avatar
Guolin Ke committed
281
282
}

283
void GBDT::UpdateScore(const Tree* tree, const int curr_class) {
Guolin Ke's avatar
Guolin Ke committed
284
  // update training score
285
  train_score_updater_->AddScore(tree_learner_.get(), curr_class);
Guolin Ke's avatar
Guolin Ke committed
286
  // update validation score
Guolin Ke's avatar
Guolin Ke committed
287
288
  for (auto& score_updater : valid_score_updater_) {
    score_updater->AddScore(tree, curr_class);
Guolin Ke's avatar
Guolin Ke committed
289
290
291
  }
}

Guolin Ke's avatar
Guolin Ke committed
292
293
294
295
std::string GBDT::OutputMetric(int iter) {
  bool need_output = (iter % gbdt_config_->output_freq) == 0;
  std::string ret = "";
  std::stringstream msg_buf;
296
  std::vector<std::pair<size_t, size_t>> meet_early_stopping_pairs;
Guolin Ke's avatar
Guolin Ke committed
297
  // print training metric
Guolin Ke's avatar
Guolin Ke committed
298
  if (need_output) {
299
300
301
    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
302
      for (size_t k = 0; k < name.size(); ++k) {
Guolin Ke's avatar
Guolin Ke committed
303
304
305
306
307
308
309
310
        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;
        }
311
      }
312
    }
Guolin Ke's avatar
Guolin Ke committed
313
314
  }
  // print validation metric
Guolin Ke's avatar
Guolin Ke committed
315
  if (need_output || early_stopping_round_ > 0) {
316
317
318
    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
319
320
321
322
323
324
325
326
327
328
329
        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;
330
          }
wxchan's avatar
wxchan committed
331
        }
Guolin Ke's avatar
Guolin Ke committed
332
        if (ret.empty() && early_stopping_round_ > 0) {
333
334
335
          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;
336
            best_iter_[i][j] = iter;
Guolin Ke's avatar
Guolin Ke committed
337
            meet_early_stopping_pairs.emplace_back(i, j);
338
          } else {
Guolin Ke's avatar
Guolin Ke committed
339
            if (iter - best_iter_[i][j] >= early_stopping_round_) { ret = best_msg_[i][j]; }
340
          }
wxchan's avatar
wxchan committed
341
342
        }
      }
Guolin Ke's avatar
Guolin Ke committed
343
344
    }
  }
Guolin Ke's avatar
Guolin Ke committed
345
346
347
  for (auto& pair : meet_early_stopping_pairs) {
    best_msg_[pair.first][pair.second] = msg_buf.str();
  }
wxchan's avatar
wxchan committed
348
  return ret;
Guolin Ke's avatar
Guolin Ke committed
349
350
}

351
/*! \brief Get eval result */
352
std::vector<double> GBDT::GetEvalAt(int data_idx) const {
Guolin Ke's avatar
Guolin Ke committed
353
  CHECK(data_idx >= 0 && data_idx <= static_cast<int>(valid_score_updater_.size()));
354
355
  std::vector<double> ret;
  if (data_idx == 0) {
356
357
    for (auto& sub_metric : training_metrics_) {
      auto scores = sub_metric->Eval(train_score_updater_->score());
358
359
360
      for (auto score : scores) {
        ret.push_back(score);
      }
361
362
    }
  }
363
364
365
366
367
368
369
  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);
      }
370
371
372
373
374
    }
  }
  return ret;
}

Guolin Ke's avatar
Guolin Ke committed
375
/*! \brief Get training scores result */
376
377
const score_t* GBDT::GetTrainingScore(int64_t* out_len) {
  *out_len = static_cast<int64_t>(train_score_updater_->num_data()) * num_class_;
Guolin Ke's avatar
Guolin Ke committed
378
  return train_score_updater_->score();
379
380
}

Guolin Ke's avatar
Guolin Ke committed
381
382
void GBDT::GetPredictAt(int data_idx, double* out_result, int64_t* out_len) {
  CHECK(data_idx >= 0 && data_idx <= static_cast<int>(valid_score_updater_.size()));
Guolin Ke's avatar
Guolin Ke committed
383
384
385
386

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

}

Guolin Ke's avatar
Guolin Ke committed
421
void GBDT::Boosting() {
422
423
424
  if (object_function_ == nullptr) {
    Log::Fatal("No object function provided");
  }
Hui Xue's avatar
Hui Xue committed
425
  // objective function will calculate gradients and hessians
426
  int64_t num_score = 0;
Guolin Ke's avatar
Guolin Ke committed
427
  object_function_->
Guolin Ke's avatar
Guolin Ke committed
428
    GetGradients(GetTrainingScore(&num_score), gradients_.data(), hessians_.data());
Guolin Ke's avatar
Guolin Ke committed
429
430
}

wxchan's avatar
wxchan committed
431
std::string GBDT::DumpModel() const {
Guolin Ke's avatar
Guolin Ke committed
432
  std::stringstream str_buf;
wxchan's avatar
wxchan committed
433

Guolin Ke's avatar
Guolin Ke committed
434
435
436
437
438
439
  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
440

Guolin Ke's avatar
Guolin Ke committed
441
442
443
444
445
446
  // 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
447
  str_buf << "\"feature_names\":[\"" 
Guolin Ke's avatar
Guolin Ke committed
448
449
450
     << Common::Join(feature_names.get(), "\",\"") << "\"]," 
     << std::endl;

Guolin Ke's avatar
Guolin Ke committed
451
  str_buf << "\"tree_info\":[";
wxchan's avatar
wxchan committed
452
453
  for (int i = 0; i < static_cast<int>(models_.size()); ++i) {
    if (i > 0) {
Guolin Ke's avatar
Guolin Ke committed
454
      str_buf << ",";
wxchan's avatar
wxchan committed
455
    }
Guolin Ke's avatar
Guolin Ke committed
456
457
458
459
    str_buf << "{";
    str_buf << "\"tree_index\":" << i << ",";
    str_buf << models_[i]->ToJSON();
    str_buf << "}";
wxchan's avatar
wxchan committed
460
  }
Guolin Ke's avatar
Guolin Ke committed
461
  str_buf << "]" << std::endl;
wxchan's avatar
wxchan committed
462

Guolin Ke's avatar
Guolin Ke committed
463
  str_buf << "}" << std::endl;
wxchan's avatar
wxchan committed
464

Guolin Ke's avatar
Guolin Ke committed
465
  return str_buf.str();
wxchan's avatar
wxchan committed
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
}

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;
483
  }
wxchan's avatar
wxchan committed
484
485
  // output sigmoid parameter
  output_file << "sigmoid=" << sigmoid_ << std::endl;
Guolin Ke's avatar
Guolin Ke committed
486
487
488
489
490
491
  // 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
492

Guolin Ke's avatar
Guolin Ke committed
493
  output_file << std::endl;
wxchan's avatar
wxchan committed
494
495
  int num_used_model = 0;
  if (num_iteration <= 0) {
Guolin Ke's avatar
Guolin Ke committed
496
497
    num_used_model = static_cast<int>(models_.size());
  } else {
wxchan's avatar
wxchan committed
498
    num_used_model = num_iteration * num_class_;
Guolin Ke's avatar
Guolin Ke committed
499
  }
wxchan's avatar
wxchan committed
500
  num_used_model = std::min(num_used_model, static_cast<int>(models_.size()));
501
  // output tree models
wxchan's avatar
wxchan committed
502
503
504
  for (int i = 0; i < num_used_model; ++i) {
    output_file << "Tree=" << i << std::endl;
    output_file << models_[i]->ToString() << std::endl;
505
  }
506

wxchan's avatar
wxchan committed
507
508
509
510
  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
511
  }
wxchan's avatar
wxchan committed
512
  output_file.close();
Guolin Ke's avatar
Guolin Ke committed
513
514
}

Guolin Ke's avatar
Guolin Ke committed
515
void GBDT::LoadModelFromString(const std::string& model_str) {
Guolin Ke's avatar
Guolin Ke committed
516
517
518
  // use serialized string to restore this object
  models_.clear();
  std::vector<std::string> lines = Common::Split(model_str.c_str(), '\n');
519
520

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

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

612
613
std::vector<double> GBDT::PredictRaw(const double* value) const {
  std::vector<double> ret(num_class_, 0.0f);
wxchan's avatar
wxchan committed
614
  for (int i = 0; i < num_iteration_for_pred_; ++i) {
615
616
617
    for (int j = 0; j < num_class_; ++j) {
      ret[j] += models_[i * num_class_ + j]->Predict(value);
    }
Guolin Ke's avatar
Guolin Ke committed
618
619
620
621
  }
  return ret;
}

622
std::vector<double> GBDT::Predict(const double* value) const {
623
  std::vector<double> ret(num_class_, 0.0f);
wxchan's avatar
wxchan committed
624
  for (int i = 0; i < num_iteration_for_pred_; ++i) {
625
626
    for (int j = 0; j < num_class_; ++j) {
      ret[j] += models_[i * num_class_ + j]->Predict(value);
627
628
    }
  }
629
630
631
632
633
634
  // 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);
  }
635
636
637
  return ret;
}

638
std::vector<int> GBDT::PredictLeafIndex(const double* value) const {
wxchan's avatar
wxchan committed
639
  std::vector<int> ret;
wxchan's avatar
wxchan committed
640
  for (int i = 0; i < num_iteration_for_pred_; ++i) {
641
642
643
    for (int j = 0; j < num_class_; ++j) {
      ret.push_back(models_[i * num_class_ + j]->PredictLeafIndex(value));
    }
wxchan's avatar
wxchan committed
644
645
646
647
  }
  return ret;
}

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