From bc1d3ee2308dddf9b583e2f0f56cff172a63a6d4 Mon Sep 17 00:00:00 2001
From: Jiaming Yuan <jm.yuan@outlook.com>
Date: Thu, 23 Jul 2020 03:28:17 +0800
Subject: [PATCH] Fix r early stop with custom objective. (#5923)

* Specify `ntreelimit`.
---
 R-package/R/utils.R                              | 5 +++--
 R-package/tests/testthat/test_custom_objective.R | 9 ++++++++-
 tests/python/test_basic_models.py                | 2 +-
 3 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/R-package/R/utils.R b/R-package/R/utils.R
index 3c4e091834a9..894e220316cf 100644
--- a/R-package/R/utils.R
+++ b/R-package/R/utils.R
@@ -145,7 +145,8 @@ xgb.iter.update <- function(booster_handle, dtrain, iter, obj = NULL) {
   if (is.null(obj)) {
     .Call(XGBoosterUpdateOneIter_R, booster_handle, as.integer(iter), dtrain)
   } else {
-    pred <- predict(booster_handle, dtrain, outputmargin = TRUE, training = TRUE)
+    pred <- predict(booster_handle, dtrain, outputmargin = TRUE, training = TRUE,
+                    ntreelimit = 0)
     gpair <- obj(pred, dtrain)
     .Call(XGBoosterBoostOneIter_R, booster_handle, dtrain, gpair$grad, gpair$hess)
   }
@@ -172,7 +173,7 @@ xgb.iter.eval <- function(booster_handle, watchlist, iter, feval = NULL) {
   } else {
     res <- sapply(seq_along(watchlist), function(j) {
       w <- watchlist[[j]]
-      preds <- predict(booster_handle, w) # predict using all trees
+      preds <- predict(booster_handle, w, ntreelimit = 0) # predict using all trees
       eval_res <- feval(preds, w)
       out <- eval_res$value
       names(out) <- paste0(evnames[j], "-", eval_res$metric)
diff --git a/R-package/tests/testthat/test_custom_objective.R b/R-package/tests/testthat/test_custom_objective.R
index 8b504e7f52e0..65bda250b0d9 100644
--- a/R-package/tests/testthat/test_custom_objective.R
+++ b/R-package/tests/testthat/test_custom_objective.R
@@ -20,7 +20,7 @@ logregobj <- function(preds, dtrain) {
 
 evalerror <- function(preds, dtrain) {
   labels <- getinfo(dtrain, "label")
-  err <- as.numeric(sum(labels != (preds > 0))) / length(labels)
+  err <- as.numeric(sum(labels != (preds > 0.5))) / length(labels)
   return(list(metric = "error", value = err))
 }
 
@@ -43,6 +43,13 @@ test_that("custom objective in CV works", {
   expect_lt(cv$evaluation_log[num_round, test_error_mean], 0.03)
 })
 
+test_that("custom objective with early stop works", {
+  bst <- xgb.train(param, dtrain, 10, watchlist)
+  expect_equal(class(bst), "xgb.Booster")
+  train_log <- bst$evaluation_log$train_error
+  expect_true(all(diff(train_log)) <= 0)
+})
+
 test_that("custom objective using DMatrix attr works", {
 
   attr(dtrain, 'label') <- getinfo(dtrain, 'label')
diff --git a/tests/python/test_basic_models.py b/tests/python/test_basic_models.py
index 9b49c5360e0c..7a1bc8c1f8f3 100644
--- a/tests/python/test_basic_models.py
+++ b/tests/python/test_basic_models.py
@@ -210,7 +210,7 @@ def logregobj(preds, dtrain):
 
         def evalerror(preds, dtrain):
             labels = dtrain.get_label()
-            return 'error', float(sum(labels != (preds > 0.0))) / len(labels)
+            return 'error', float(sum(labels != (preds > 0.5))) / len(labels)
 
         # test custom_objective in training
         bst = xgb.train(param, dtrain, num_round, watchlist, logregobj, evalerror)