[XGBoost4J-Spark] bestIteration and bestScore for early stopping (#7095)

dmlc · Jul 19, 2021 · 9f7f8b9 · 9f7f8b9
1 parent d7c1449
commit 9f7f8b9
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/jvm-packages/xgboost4j/src/main/java/ml/dmlc/xgboost4j/java/XGBoost.java b/jvm-packages/xgboost4j/src/main/java/ml/dmlc/xgboost4j/java/XGBoost.java
@@ -242,11 +242,15 @@ public static Booster trainAndSaveCheckpoint(
           if (score > bestScore) {
             bestScore = score;
             bestIteration = iter;
+            booster.setAttr("best_iteration", String.valueOf(bestIteration));
+            booster.setAttr("best_score", String.valueOf(bestScore));
           }
         } else {
           if (score < bestScore) {
             bestScore = score;
             bestIteration = iter;
+            booster.setAttr("best_iteration", String.valueOf(bestIteration));
+            booster.setAttr("best_score", String.valueOf(bestScore));
           }
         }
         if (earlyStoppingRounds > 0) {

diff --git a/jvm-packages/xgboost4j/src/test/java/ml/dmlc/xgboost4j/java/BoosterImplTest.java b/jvm-packages/xgboost4j/src/test/java/ml/dmlc/xgboost4j/java/BoosterImplTest.java
@@ -643,7 +643,7 @@ public void testSetAndGetAttrs() throws XGBoostError {
     }});
 
     Map<String, String> attr = booster.getAttrs();
-    TestCase.assertEquals(attr.size(), 4);
+    TestCase.assertEquals(attr.size(), 6);
     TestCase.assertEquals(attr.get("testKey1"), "testValue2");
     TestCase.assertEquals(attr.get("aa"), "AA");
     TestCase.assertEquals(attr.get("bb"), "BB");