Commit 67f98d66 authored by Guolin Ke's avatar Guolin Ke
Browse files

output the num_leaves and max_depth after training a tree. (#303)

parent 8b00b8e1
...@@ -146,6 +146,7 @@ Tree* SerialTreeLearner::Train(const score_t* gradients, const score_t *hessians ...@@ -146,6 +146,7 @@ Tree* SerialTreeLearner::Train(const score_t* gradients, const score_t *hessians
last_trained_tree_ = tree.get(); last_trained_tree_ = tree.get();
// root leaf // root leaf
int left_leaf = 0; int left_leaf = 0;
int cur_depth = 1;
// only root leaf can be splitted on first time // only root leaf can be splitted on first time
int right_leaf = -1; int right_leaf = -1;
for (int split = 0; split < tree_config_->num_leaves - 1; ++split) { for (int split = 0; split < tree_config_->num_leaves - 1; ++split) {
...@@ -162,13 +163,14 @@ Tree* SerialTreeLearner::Train(const score_t* gradients, const score_t *hessians ...@@ -162,13 +163,14 @@ Tree* SerialTreeLearner::Train(const score_t* gradients, const score_t *hessians
const SplitInfo& best_leaf_SplitInfo = best_split_per_leaf_[best_leaf]; const SplitInfo& best_leaf_SplitInfo = best_split_per_leaf_[best_leaf];
// cannot split, quit // cannot split, quit
if (best_leaf_SplitInfo.gain <= 0.0) { if (best_leaf_SplitInfo.gain <= 0.0) {
Log::Info("No further splits with positive gain, best gain: %f, leaves: %d", Log::Info("No further splits with positive gain, best gain: %f", best_leaf_SplitInfo.gain);
best_leaf_SplitInfo.gain, split + 1);
break; break;
} }
// split tree with best leaf // split tree with best leaf
Split(tree.get(), best_leaf, &left_leaf, &right_leaf); Split(tree.get(), best_leaf, &left_leaf, &right_leaf);
cur_depth = std::max(cur_depth, tree->leaf_depth(left_leaf));
} }
Log::Info("Trained a tree with leaves=%d and max_depth=%d", tree->num_leaves(), cur_depth);
return tree.release(); return tree.release();
} }
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment