store the true split gain in tree model (#3196)

cfc5e4fe · shiyu1994 · GitHub · ddf8c104 · cfc5e4fe
Unverified Commit cfc5e4fe authored Jul 02, 2020 by shiyu1994 Committed by GitHub Jul 02, 2020
Hide whitespace changes
Inline Side-by-side

Showing with 4 additions and 2 deletions

src/treelearner/serial_tree_learner.cpp src/treelearner/serial_tree_learner.cpp +4 -2

No files found.
--- a/src/treelearner/serial_tree_learner.cpp
+++ b/src/treelearner/serial_tree_learner.cpp
@@ -577,7 +577,8 @@ void SerialTreeLearner::SplitInner(Tree* tree, int best_leaf, int* left_leaf,
        static_cast<data_size_t>(best_split_info.right_count),
        static_cast<double>(best_split_info.left_sum_hessian),
        static_cast<double>(best_split_info.right_sum_hessian),
-        static_cast<float>(best_split_info.gain),
+        // store the true split gain in tree model
+        static_cast<float>(best_split_info.gain + config_->min_gain_to_split),
        train_data_->FeatureBinMapper(inner_feature_index)->missing_type(),
        best_split_info.default_left);
  } else {
@@ -613,7 +614,8 @@ void SerialTreeLearner::SplitInner(Tree* tree, int best_leaf, int* left_leaf,
        static_cast<data_size_t>(best_split_info.right_count),
        static_cast<double>(best_split_info.left_sum_hessian),
        static_cast<double>(best_split_info.right_sum_hessian),
-        static_cast<float>(best_split_info.gain),
+        // store the true split gain in tree model
+        static_cast<float>(best_split_info.gain + config_->min_gain_to_split),
        train_data_->FeatureBinMapper(inner_feature_index)->missing_type());
  }