From f55e37d02670374085d523357f0032ac544b6ed2 Mon Sep 17 00:00:00 2001
From: Guolin Ke <guolin.ke@outlook.com>
Date: Sat, 1 Dec 2018 12:16:46 +0800
Subject: [PATCH] [R] add reset_data in lgb.cv (#1874)

* add free booster in lgb.cv

* change to reset_data
---
 R-package/R/lgb.cv.R    | 16 +++++++++++++++-
 R-package/R/lgb.train.R |  2 ++
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/R-package/R/lgb.cv.R b/R-package/R/lgb.cv.R
index ddaf5c281998..c0d4bc2e6335 100644
--- a/R-package/R/lgb.cv.R
+++ b/R-package/R/lgb.cv.R
@@ -41,6 +41,7 @@ CVBooster <- R6::R6Class(
 #'        type str represents feature names
 #' @param callbacks list of callback functions
 #'        List of callback functions that are applied at each iteration.
+#' @param reset_data Boolean, setting it to TRUE (not the default value) will transform the booster model into a predictor model which frees up memory and the original datasets
 #' @param ... other parameters, see Parameters.rst for more information. A few key parameters:
 #'            \itemize{
 #'                \item{boosting}{Boosting type. \code{"gbdt"} or \code{"dart"}}
@@ -87,6 +88,7 @@ lgb.cv <- function(params = list(),
                    categorical_feature = NULL,
                    early_stopping_rounds = NULL,
                    callbacks = list(),
+                   reset_data = FALSE,
                    ...) {
 
   # Setup temporary variables
@@ -303,7 +305,19 @@ lgb.cv <- function(params = list(),
     if (env$met_early_stop) break
 
   }
-
+  if (reset_data) {
+    lapply(cv_booster$boosters, function(fd) {
+      # Store temporarily model data elsewhere
+      booster_old <- list(best_iter = fd$booster$best_iter,
+                          best_score = fd$booster$best_score,
+                          record_evals = fd$booster$record_evals)
+      # Reload model
+      fd$booster <- lgb.load(model_str = fd$booster$save_model_to_string())
+      fd$booster$best_iter <- booster_old$best_iter
+      fd$booster$best_score <- booster_old$best_score
+      fd$booster$record_evals <- booster_old$record_evals
+    })
+  }
   # Return booster
   return(cv_booster)
 
diff --git a/R-package/R/lgb.train.R b/R-package/R/lgb.train.R
index 29a37853e899..f9744e8d3ce1 100644
--- a/R-package/R/lgb.train.R
+++ b/R-package/R/lgb.train.R
@@ -12,6 +12,8 @@
 #' @param categorical_feature list of str or int
 #'        type int represents index,
 #'        type str represents feature names
+#' @param callbacks list of callback functions
+#'        List of callback functions that are applied at each iteration.
 #' @param reset_data Boolean, setting it to TRUE (not the default value) will transform the booster model into a predictor model which frees up memory and the original datasets
 #' @param ... other parameters, see Parameters.rst for more information. A few key parameters:
 #'            \itemize{