From e94b76631035cd8b3a5cdd0c883225f069e74686 Mon Sep 17 00:00:00 2001
From: Jiaming Yuan <jm.yuan@outlook.com>
Date: Thu, 13 Jan 2022 21:53:06 +0800
Subject: [PATCH] Fix early stopping with linear model. (#7554)

---
 R-package/tests/testthat/test_glm.R | 28 ++++++++++++++++++++++++++++
 src/gbm/gblinear.cc                 | 11 +++++------
 2 files changed, 33 insertions(+), 6 deletions(-)

diff --git a/R-package/tests/testthat/test_glm.R b/R-package/tests/testthat/test_glm.R
index 99e73968d7ce..270267a0fe0c 100644
--- a/R-package/tests/testthat/test_glm.R
+++ b/R-package/tests/testthat/test_glm.R
@@ -46,3 +46,31 @@ test_that("gblinear works", {
   expect_equal(dim(h), c(n, ncol(dtrain) + 1))
   expect_s4_class(h, "dgCMatrix")
 })
+
+test_that("gblinear early stopping works", {
+  data(agaricus.train, package = 'xgboost')
+  data(agaricus.test, package = 'xgboost')
+  dtrain <- xgb.DMatrix(agaricus.train$data, label = agaricus.train$label)
+  dtest <- xgb.DMatrix(agaricus.test$data, label = agaricus.test$label)
+
+  param <- list(
+    objective = "binary:logistic", eval_metric = "error", booster = "gblinear",
+    nthread = 2, eta = 0.8, alpha = 0.0001, lambda = 0.0001,
+    updater = "coord_descent"
+  )
+
+  es_round <- 1
+  n <- 10
+  booster <- xgb.train(
+    param, dtrain, n, list(eval = dtest, train = dtrain), early_stopping_rounds = es_round
+  )
+  expect_equal(booster$best_iteration, 5)
+  predt_es <- predict(booster, dtrain)
+
+  n <- booster$best_iteration + es_round
+  booster <- xgb.train(
+    param, dtrain, n, list(eval = dtest, train = dtrain), early_stopping_rounds = es_round
+  )
+  predt <- predict(booster, dtrain)
+  expect_equal(predt_es, predt)
+})
diff --git a/src/gbm/gblinear.cc b/src/gbm/gblinear.cc
index 31f0d632fe26..a1348bb4d0ce 100644
--- a/src/gbm/gblinear.cc
+++ b/src/gbm/gblinear.cc
@@ -62,9 +62,8 @@ struct GBLinearTrainParam : public XGBoostParameter<GBLinearTrainParam> {
   }
 };
 
-void LinearCheckLayer(unsigned layer_begin, unsigned layer_end) {
+void LinearCheckLayer(unsigned layer_begin) {
   CHECK_EQ(layer_begin, 0) << "Linear booster does not support prediction range.";
-  CHECK_EQ(layer_end, 0)   << "Linear booster does not support prediction range.";
 }
 
 /*!
@@ -152,7 +151,7 @@ class GBLinear : public GradientBooster {
   void PredictBatch(DMatrix *p_fmat, PredictionCacheEntry *predts,
                     bool training, unsigned layer_begin, unsigned layer_end) override {
     monitor_.Start("PredictBatch");
-    LinearCheckLayer(layer_begin, layer_end);
+    LinearCheckLayer(layer_begin);
     auto* out_preds = &predts->predictions;
     this->PredictBatchInternal(p_fmat, &out_preds->HostVector());
     monitor_.Stop("PredictBatch");
@@ -161,7 +160,7 @@ class GBLinear : public GradientBooster {
   void PredictInstance(const SparsePage::Inst &inst,
                        std::vector<bst_float> *out_preds,
                        unsigned layer_begin, unsigned layer_end) override {
-    LinearCheckLayer(layer_begin, layer_end);
+    LinearCheckLayer(layer_begin);
     const int ngroup = model_.learner_model_param->num_output_group;
     for (int gid = 0; gid < ngroup; ++gid) {
       this->Pred(inst, dmlc::BeginPtr(*out_preds), gid,
@@ -177,7 +176,7 @@ class GBLinear : public GradientBooster {
                            HostDeviceVector<bst_float>* out_contribs,
                            unsigned layer_begin, unsigned layer_end, bool, int, unsigned) override {
     model_.LazyInitModel();
-    LinearCheckLayer(layer_begin, layer_end);
+    LinearCheckLayer(layer_begin);
     auto base_margin = p_fmat->Info().base_margin_.View(GenericParameter::kCpuId);
     const int ngroup = model_.learner_model_param->num_output_group;
     const size_t ncolumns = model_.learner_model_param->num_feature + 1;
@@ -214,7 +213,7 @@ class GBLinear : public GradientBooster {
   void PredictInteractionContributions(DMatrix* p_fmat,
                                        HostDeviceVector<bst_float>* out_contribs,
                                        unsigned layer_begin, unsigned layer_end, bool) override {
-    LinearCheckLayer(layer_begin, layer_end);
+    LinearCheckLayer(layer_begin);
     std::vector<bst_float>& contribs = out_contribs->HostVector();
 
     // linear models have no interaction effects