fix tests.

Alnusjaponica · Mar 22, 2017 · 2e962c7 · 2e962c7
1 parent e179c7c
commit 2e962c7
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 5 deletions.
diff --git a/src/boosting/gbdt.cpp b/src/boosting/gbdt.cpp
@@ -630,6 +630,7 @@ std::string GBDT::DumpModel(int num_iteration) const {
   str_buf << "\"tree_info\":[";
   int num_used_model = static_cast<int>(models_.size());
   if (num_iteration > 0) {
+    num_iteration += boost_from_average_ ? 1 : 0;
     num_used_model = std::min(num_iteration * num_class_, num_used_model);
   }
   for (int i = 0; i < num_used_model; ++i) {
@@ -648,7 +649,7 @@ std::string GBDT::DumpModel(int num_iteration) const {
   return str_buf.str();
 }
 
-std::string GBDT::SaveModelToString(int num_iterations) const {
+std::string GBDT::SaveModelToString(int num_iteration) const {
   std::stringstream ss;
 
   // output model type
@@ -676,8 +677,9 @@ std::string GBDT::SaveModelToString(int num_iterations) const {
 
   ss << std::endl;
   int num_used_model = static_cast<int>(models_.size());
-  if (num_iterations > 0) {
-    num_used_model = std::min(num_iterations * num_class_, num_used_model);
+  if (num_iteration > 0) {
+    num_iteration += boost_from_average_ ? 1 : 0;
+    num_used_model = std::min(num_iteration * num_class_, num_used_model);
   }
   // output tree models
   for (int i = 0; i < num_used_model; ++i) {

diff --git a/src/boosting/gbdt.h b/src/boosting/gbdt.h
@@ -89,7 +89,7 @@ class GBDT: public Boosting {
   */
   void RollbackOneIter() override;
 
-  int GetCurrentIteration() const override { return iter_ + num_init_iteration_; }
+  int GetCurrentIteration() const override { return static_cast<int>(models_.size()) / num_class_; }
 
   bool EvalAndCheckEarlyStopping() override;
 

diff --git a/tests/python_package_test/test_engine.py b/tests/python_package_test/test_engine.py
@@ -32,7 +32,7 @@ class template(object):
     @staticmethod
     def test_template(params={'objective': 'regression', 'metric': 'l2'},
                       X_y=load_boston(True), feval=mean_squared_error,
-                      num_round=150, init_model=None, custom_eval=None,
+                      num_round=200, init_model=None, custom_eval=None,
                       early_stopping_rounds=10,
                       return_data=False, return_model=False):
         params['verbose'], params['seed'] = -1, 42