Merge pull request #304 from tidymodels/early-stop

Early stopping in xgboost

Author: topepo

DESCRIPTION

@@ -26,7 +26,7 @@ Imports:
     glue,
     magrittr,
     stats,
-    tidyr,
+    tidyr (>= 1.0.0),
     globals,
     prettyunits,
     vctrs (>= 0.2.0)

NEWS.md

@@ -1,5 +1,13 @@
 # parsnip (development version)
 
+## Other Changes
+
+ * `tidyr` >= 1.0.0 is now required. 
+ 
+## New Features
+
+ * A new main argument was added to `boost_tree()` called `stop_iter` for early stopping. The `xgb_train()` function gained arguments for early stopping and a percentage of data to leave out for a validation set. 
+ 
 # parsnip 0.1.1
 
 ## New Features

R/boost_tree.R

@@ -19,6 +19,8 @@
 #'   \item \code{loss_reduction}: The reduction in the loss function required
 #'   to split further.
 #'   \item \code{sample_size}: The amount of data exposed to the fitting routine.
+#'   \item \code{stop_iter}: The number of iterations without improvement before
+#'   stopping.
 #' }
 #' These arguments are converted to their specific names at the
 #'  time that the model is fit. Other options and argument can be
@@ -46,6 +48,8 @@
 #' @param sample_size A number for the number (or proportion) of data that is
 #'  exposed to the fitting routine. For `xgboost`, the sampling is done at at
 #'  each iteration while `C5.0` samples once during training.
+#' @param stop_iter The number of iterations without improvement before
+#'   stopping  (`xgboost` only).
 #' @details
 #' The data given to the function are not saved and are only used
 #'  to determine the _mode_ of the model. For `boost_tree()`, the
@@ -87,15 +91,17 @@ boost_tree <-
            mtry = NULL, trees = NULL, min_n = NULL,
            tree_depth = NULL, learn_rate = NULL,
            loss_reduction = NULL,
-           sample_size = NULL) {
+           sample_size = NULL,
+           stop_iter = NULL) {
     args <- list(
       mtry = enquo(mtry),
       trees = enquo(trees),
       min_n = enquo(min_n),
       tree_depth = enquo(tree_depth),
       learn_rate = enquo(learn_rate),
       loss_reduction = enquo(loss_reduction),
-      sample_size = enquo(sample_size)
+      sample_size = enquo(sample_size),
+      stop_iter = enquo(stop_iter)
     )
 
     new_model_spec(
@@ -155,6 +161,7 @@ update.boost_tree <-
            mtry = NULL, trees = NULL, min_n = NULL,
            tree_depth = NULL, learn_rate = NULL,
            loss_reduction = NULL, sample_size = NULL,
+           stop_iter = NULL,
            fresh = FALSE, ...) {
     update_dot_check(...)
 
@@ -169,7 +176,8 @@ update.boost_tree <-
       tree_depth = enquo(tree_depth),
       learn_rate = enquo(learn_rate),
       loss_reduction = enquo(loss_reduction),
-      sample_size = enquo(sample_size)
+      sample_size = enquo(sample_size),
+      stop_iter = enquo(stop_iter)
     )
 
     args <- update_main_parameters(args, parameters)
@@ -242,8 +250,8 @@ check_args.boost_tree <- function(object) {
 
 #' Boosted trees via xgboost
 #'
-#' `xgb_train` is a wrapper for `xgboost` tree-based models
-#'  where all of the model arguments are in the main function.
+#' `xgb_train` is a wrapper for `xgboost` tree-based models where all of the
+#'  model arguments are in the main function.
 #'
 #' @param x A data frame or matrix of predictors
 #' @param y A vector (factor or numeric) or matrix (numeric) of outcome data.
@@ -256,16 +264,41 @@ check_args.boost_tree <- function(object) {
 #' @param gamma A number for the minimum loss reduction required to make a
 #'  further partition on a leaf node of the tree
 #' @param subsample Subsampling proportion of rows.
+#' @param validation A positive number. If on `[0, 1)` the value, `validation`
+#' is a random proportion of data in `x` and `y` that are used for performance
+#' assessment and potential early stopping. If 1 or greater, it is the _number_
+#' of training set samples use for these purposes.
+#' @param early_stop An integer or `NULL`. If not `NULL`, it is the number of
+#' training iterations without improvement before stopping. If `validation` is
+#' used, performance is base on the validation set; otherwise the training set
+#' is used.
 #' @param ... Other options to pass to `xgb.train`.
 #' @return A fitted `xgboost` object.
 #' @keywords internal
 #' @export
 xgb_train <- function(
   x, y,
   max_depth = 6, nrounds = 15, eta  = 0.3, colsample_bytree = 1,
-  min_child_weight = 1, gamma = 0, subsample = 1, ...) {
+  min_child_weight = 1, gamma = 0, subsample = 1, validation = 0,
+  early_stop = NULL, ...) {
+
+  if (length(levels(y)) > 2) {
+    num_class <- length(levels(y))
+  }  else {
+    num_class <- NULL
+  }
+  if (!is.numeric(validation) || validation < 0 || validation >= 1) {
+    rlang::abort("`validation` should be on [0, 1).")
+  }
+  if (!is.null(early_stop)) {
+    if (early_stop <= 1) {
+      rlang::abort(paste0("`early_stop` should be on [2, ",  nrounds, ")."))
+    } else if (early_stop >= nrounds) {
+      early_stop <- nrounds - 1
+      rlang::warn(paste0("`early_stop` was reduced to ", early_stop, "."))
+    }
+  }
 
-  num_class <- if (length(levels(y)) > 2) length(levels(y)) else NULL
 
   if (is.numeric(y)) {
     loss <- "reg:linear"
@@ -287,7 +320,16 @@ xgb_train <- function(
   p <- ncol(x)
 
   if (!inherits(x, "xgb.DMatrix")) {
-    x <- xgboost::xgb.DMatrix(x, label = y, missing = NA)
+    if (validation > 0) {
+      trn_index <- sample(1:n, size = floor(n * validation) + 1)
+      wlist <-
+        list(validation = xgboost::xgb.DMatrix(x[-trn_index, ], label = y[-trn_index], missing = NA))
+      x <- xgboost::xgb.DMatrix(x[trn_index, ], label = y[trn_index], missing = NA)
+
+    } else {
+      x <- xgboost::xgb.DMatrix(x, label = y, missing = NA)
+      wlist <- list(training = x)
+    }
   } else {
     xgboost::setinfo(x, "label", y)
   }
@@ -320,9 +362,11 @@ xgb_train <- function(
 
   main_args <- list(
     data = quote(x),
+    watchlist = quote(wlist),
     params = arg_list,
     nrounds = nrounds,
-    objective = loss
+    objective = loss,
+    early_stopping_rounds = early_stop
   )
   if (!is.null(num_class)) {
     main_args$num_class <- num_class
@@ -334,6 +378,9 @@ xgb_train <- function(
   others <- list(...)
   others <-
     others[!(names(others) %in% c("data", "weights", "nrounds", "num_class", names(arg_list)))]
+  if (!(any(names(others) == "verbose"))) {
+    others$verbose <- 0
+  }
   if (length(others) > 0) {
     call <- rlang::call_modify(call, !!!others)
   }

R/boost_tree_data.R

@@ -65,6 +65,15 @@ set_model_arg(
   func = list(pkg = "dials", fun = "sample_size"),
   has_submodel = FALSE
 )
+set_model_arg(
+  model = "boost_tree",
+  eng = "xgboost",
+  parsnip = "stop_iter",
+  original = "early_stop",
+  func = list(pkg = "dials", fun = "stop_iter"),
+  has_submodel = FALSE
+)
+
 
 set_fit(
   model = "boost_tree",

_pkgdown.yml

@@ -1,7 +1,7 @@
 template:
   package: tidytemplate
   params:
-    part_of: <a href="https://github.com/tidymodels">tidymodels</a>
+    part_of: <a href="https://tidymodels.org">tidymodels</a>
     footer: <code>parsnip</code> is a part of the <strong>tidymodels</strong> ecosystem, a collection of modeling packages designed with common APIs and a shared philosophy.
 
 # https://github.com/tidyverse/tidytemplate for css