R/sgd.R

In cuda.ml: R Interface for the RAPIDS cuML Suite of Libraries

sgd_match_penalty <- function(penalty = c("none", "l1", "l2", "elasticnet")) {
  penalty <- match.arg(penalty)

  switch(penalty, none = 0L, l1 = 1L, l2 = 2L, elasticnet = 3L)
}

sgd_match_learning_rate <- function(
  learning_rate = c("constant", "invscaling", "adaptive")
) {
  learning_rate <- match.arg(learning_rate)

  switch(learning_rate, constant = 1L, invscaling = 2L, adaptive = 3L)
}

#' Train a linear model using mini-batch stochastic gradient descent.
#'
#' Train a linear model using mini-batch stochastic gradient descent.
#'
#' @template supervised-model-inputs
#' @template supervised-model-output
#' @template ellipsis-unused
#' @template fit-intercept
#' @template l1_ratio
#' @param penalty Type of regularization to perform, must be one of
#'   \{"none", "l1", "l2", "elasticnet"\}.
#'
#'   - "none": no regularization.
#'   - "l1": perform regularization based on the L1-norm (LASSO) which tries to
#'           minimize the sum of the absolute values of the coefficients.
#'   - "l2": perform regularization based on the L2 norm (Ridge) which tries to
#'           minimize the sum of the square of the coefficients.
#'   - "elasticnet": perform the Elastic Net regularization which is based on
#'                   the weighted average of L1 and L2 norms.
#'
#'   Default: "none".
#' @param alpha Multiplier of the penalty term. Default: 1e-4.
#' @param batch_size The number of samples that will be included in each batch.
#'   Default: 32L.
#' @param epochs The number of times the model should iterate through the entire
#'   dataset during training. Default: 1000L.
#' @param tol Threshold for stopping training. Training will stop if
#'   (loss in current epoch) > (loss in previous epoch) - \code{tol}.
#'   Default: 1e-3.
#' @param shuffle Whether to shuffle the training data after each epoch.
#'   Default: TRUE.
#' @param eta0 The initial learning rate. Default: 1e-3.
#' @param power_t The exponent used for calculating the invscaling learning
#'   rate. Default: 0.5.
#' @param learning_rate Must be one of \{"constant", "invscaling", "adaptive"\}.
#'
#'   - "constant": the learning rate will be kept constant.
#'   - "invscaling": (learning rate) = (initial learning rate) / pow(t, power_t)
#'                   where \code{t} is the number of epochs and
#'                   \code{power_t} is a tunable parameter of this model.
#'   - "adaptive": (learning rate) = (initial learning rate) as long as the
#'                 training loss keeps decreasing. Each time the last
#'                 \code{n_iter_no_change} consecutive epochs fail to decrease
#'                 the training loss by \code{tol}, the current learning rate is
#'                 divided by 5.
#'
#'   Default: "constant".
#' @param n_iter_no_change The maximum number of epochs to train if there is no
#'   improvement in the model. Default: 5.
#'
#' @return A linear model that can be used with the 'predict' S3 generic to make
#'   predictions on new data points.
#'
#' @examples
#'
#' library(cuda.ml)
#'
#' if (interactive() && cuda_ml_backend_info()$runtime_installed) {
#'   model <- cuda_ml_sgd(
#'     mpg ~ ., mtcars,
#'     batch_size = 4, epochs = 50000,
#'     learning_rate = "adaptive", eta0 = 1e-5,
#'     penalty = "l2", alpha = 1e-5, tol = 1e-6,
#'     n_iter_no_change = 10
#'   )
#'
#'   predictors <- subset(mtcars, select = -mpg)
#'   preds <- predict(model, predictors)
#'   print(all.equal(preds$.pred, mtcars$mpg, tolerance = 0.09))
#' }
#' @importFrom ellipsis check_dots_used
#' @export
cuda_ml_sgd <- function(x, ...) {
  UseMethod("cuda_ml_sgd")
}

#' @rdname cuda_ml_sgd
#' @export
cuda_ml_sgd.default <- function(x, ...) {
  report_undefined_fn("cuda_ml_sgd", x)
}

#' @rdname cuda_ml_sgd
#' @export
cuda_ml_sgd.data.frame <- function(
  x,
  y,
  fit_intercept = TRUE,
  penalty = c("none", "l1", "l2", "elasticnet"),
  alpha = 1e-4,
  l1_ratio = 0.5,
  epochs = 1000L,
  tol = 1e-3,
  shuffle = TRUE,
  learning_rate = c("constant", "invscaling", "adaptive"),
  eta0 = 1e-3,
  power_t = 0.5,
  batch_size = 32L,
  n_iter_no_change = 5L,
  ...
) {
  check_dots_used()
  processed <- hardhat::mold(x, y)

  cuda_ml_sgd_bridge(
    processed = processed,
    fit_intercept = fit_intercept,
    penalty = penalty,
    alpha = alpha,
    l1_ratio = l1_ratio,
    epochs = epochs,
    tol = tol,
    shuffle = shuffle,
    learning_rate = learning_rate,
    eta0 = eta0,
    power_t = power_t,
    batch_size = batch_size,
    n_iter_no_change = n_iter_no_change
  )
}

#' @rdname cuda_ml_sgd
#' @export
cuda_ml_sgd.matrix <- function(
  x,
  y,
  fit_intercept = TRUE,
  penalty = c("none", "l1", "l2", "elasticnet"),
  alpha = 1e-4,
  l1_ratio = 0.5,
  epochs = 1000L,
  tol = 1e-3,
  shuffle = TRUE,
  learning_rate = c("constant", "invscaling", "adaptive"),
  eta0 = 1e-3,
  power_t = 0.5,
  batch_size = 32L,
  n_iter_no_change = 5L,
  ...
) {
  check_dots_used()
  processed <- hardhat::mold(x, y)

  cuda_ml_sgd_bridge(
    processed = processed,
    fit_intercept = fit_intercept,
    penalty = penalty,
    alpha = alpha,
    l1_ratio = l1_ratio,
    epochs = epochs,
    tol = tol,
    shuffle = shuffle,
    learning_rate = learning_rate,
    eta0 = eta0,
    power_t = power_t,
    batch_size = batch_size,
    n_iter_no_change = n_iter_no_change
  )
}

#' @rdname cuda_ml_sgd
#' @export
cuda_ml_sgd.formula <- function(
  formula,
  data,
  fit_intercept = TRUE,
  penalty = c("none", "l1", "l2", "elasticnet"),
  alpha = 1e-4,
  l1_ratio = 0.5,
  epochs = 1000L,
  tol = 1e-3,
  shuffle = TRUE,
  learning_rate = c("constant", "invscaling", "adaptive"),
  eta0 = 1e-3,
  power_t = 0.5,
  batch_size = 32L,
  n_iter_no_change = 5L,
  ...
) {
  check_dots_used()
  processed <- hardhat::mold(formula, data)

  cuda_ml_sgd_bridge(
    processed = processed,
    fit_intercept = fit_intercept,
    penalty = penalty,
    alpha = alpha,
    l1_ratio = l1_ratio,
    epochs = epochs,
    tol = tol,
    shuffle = shuffle,
    learning_rate = learning_rate,
    eta0 = eta0,
    power_t = power_t,
    batch_size = batch_size,
    n_iter_no_change = n_iter_no_change
  )
}

#' @rdname cuda_ml_sgd
#' @export
cuda_ml_sgd.recipe <- function(
  x,
  data,
  fit_intercept = TRUE,
  penalty = c("none", "l1", "l2", "elasticnet"),
  alpha = 1e-4,
  l1_ratio = 0.5,
  epochs = 1000L,
  tol = 1e-3,
  shuffle = TRUE,
  learning_rate = c("constant", "invscaling", "adaptive"),
  eta0 = 1e-3,
  power_t = 0.5,
  batch_size = 32L,
  n_iter_no_change = 5L,
  ...
) {
  check_dots_used()
  processed <- hardhat::mold(x, data)

  cuda_ml_sgd_bridge(
    processed = processed,
    fit_intercept = fit_intercept,
    penalty = penalty,
    alpha = alpha,
    l1_ratio = l1_ratio,
    epochs = epochs,
    tol = tol,
    shuffle = shuffle,
    learning_rate = learning_rate,
    eta0 = eta0,
    power_t = power_t,
    batch_size = batch_size,
    n_iter_no_change = n_iter_no_change
  )
}

cuda_ml_sgd_bridge <- function(
  processed,
  fit_intercept,
  penalty,
  alpha,
  l1_ratio,
  epochs,
  tol,
  shuffle,
  learning_rate,
  eta0,
  power_t,
  batch_size,
  n_iter_no_change
) {
  validate_lm_input(processed)
  penalty <- sgd_match_penalty(penalty)

  learning_rate <- sgd_match_learning_rate(learning_rate)
  if (eta0 <= 0) {
    stop("eta0 must be > 0")
  }

  x <- as.matrix(processed$predictors)
  y <- processed$outcomes[[1]]

  model_xptr <- .sgd_fit(
    x = x,
    y = y,
    fit_intercept = fit_intercept,
    batch_size = as.integer(batch_size),
    epochs = as.integer(epochs),
    lr_type = learning_rate,
    eta0 = as.numeric(eta0),
    power_t = as.numeric(power_t),
    loss = 0L,
    penalty = penalty,
    alpha = as.numeric(alpha),
    l1_ratio = as.numeric(l1_ratio),
    shuffle = shuffle,
    tol = as.numeric(tol),
    n_iter_no_change = as.integer(n_iter_no_change)
  )

  new_linear_model(
    cls = "cuda_ml_sgd",
    xptr = model_xptr,
    blueprint = processed$blueprint
  )
}