lightgbm.R 2.95 KB
Newer Older
1
2
3
#' Simple interface for training an lightgbm model.
#' Its documentation is combined with lgb.train.
#'
Guolin Ke's avatar
Guolin Ke committed
4
5
#' @rdname lgb.train
#' @export
6
7
8
9
10
11
12
13
14
15
16
17
18
19
lightgbm <- function(data,
                     label = NULL,
                     weight = NULL,
                     params = list(),
                     nrounds = 10,
                     verbose = 1,
                     eval_freq = 1L,
                     early_stopping_rounds = NULL,
                     save_name = "lightgbm.model",
                     init_model = NULL,
                     callbacks = list(),
                     ...) {
  
  # Set data to a temporary variable
Guolin Ke's avatar
Guolin Ke committed
20
  dtrain <- data
21
22
  
  # Check whether data is lgb.Dataset, if not then create lgb.Dataset manually
23
24
  if (!lgb.is.Dataset(dtrain)) {
    dtrain <- lgb.Dataset(data, label = label, weight = weight)
Guolin Ke's avatar
Guolin Ke committed
25
  }
Guolin Ke's avatar
Guolin Ke committed
26

27
  # Set validation as oneself
Guolin Ke's avatar
Guolin Ke committed
28
  valids <- list()
29
30
31
32
33
  if (verbose > 0) {
    valids$train = dtrain
  }
  
  # Train a model using the regular way
34
  bst <- lgb.train(params, dtrain, nrounds, valids, verbose = verbose, eval_freq = eval_freq,
Guolin Ke's avatar
Guolin Ke committed
35
36
                   early_stopping_rounds = early_stopping_rounds,
                   init_model = init_model, callbacks = callbacks, ...)
37
38
  
  # Store model under a specific name
Guolin Ke's avatar
Guolin Ke committed
39
  bst$save_model(save_name)
40
41
42
  
  # Return booster
  return(bst)
Guolin Ke's avatar
Guolin Ke committed
43
44
45
}

#' Training part from Mushroom Data Set
46
#'
Guolin Ke's avatar
Guolin Ke committed
47
48
#' This data set is originally from the Mushroom data set,
#' UCI Machine Learning Repository.
49
#'
Guolin Ke's avatar
Guolin Ke committed
50
#' This data set includes the following fields:
51
#'
Guolin Ke's avatar
Guolin Ke committed
52
53
54
55
56
57
58
#' \itemize{
#'  \item \code{label} the label for each record
#'  \item \code{data} a sparse Matrix of \code{dgCMatrix} class, with 126 columns.
#' }
#'
#' @references
#' https://archive.ics.uci.edu/ml/datasets/Mushroom
59
60
61
#'
#' Bache, K. & Lichman, M. (2013). UCI Machine Learning Repository
#' [http://archive.ics.uci.edu/ml]. Irvine, CA: University of California,
Guolin Ke's avatar
Guolin Ke committed
62
#' School of Information and Computer Science.
63
#'
Guolin Ke's avatar
Guolin Ke committed
64
65
66
67
#' @docType data
#' @keywords datasets
#' @name agaricus.train
#' @usage data(agaricus.train)
68
#' @format A list containing a label vector, and a dgCMatrix object with 6513
Guolin Ke's avatar
Guolin Ke committed
69
70
71
72
73
74
75
#' rows and 127 variables
NULL

#' Test part from Mushroom Data Set
#'
#' This data set is originally from the Mushroom data set,
#' UCI Machine Learning Repository.
76
#'
Guolin Ke's avatar
Guolin Ke committed
77
#' This data set includes the following fields:
78
#'
Guolin Ke's avatar
Guolin Ke committed
79
80
81
82
83
84
85
#' \itemize{
#'  \item \code{label} the label for each record
#'  \item \code{data} a sparse Matrix of \code{dgCMatrix} class, with 126 columns.
#' }
#'
#' @references
#' https://archive.ics.uci.edu/ml/datasets/Mushroom
86
87
88
#'
#' Bache, K. & Lichman, M. (2013). UCI Machine Learning Repository
#' [http://archive.ics.uci.edu/ml]. Irvine, CA: University of California,
Guolin Ke's avatar
Guolin Ke committed
89
#' School of Information and Computer Science.
90
#'
Guolin Ke's avatar
Guolin Ke committed
91
92
93
94
#' @docType data
#' @keywords datasets
#' @name agaricus.test
#' @usage data(agaricus.test)
95
#' @format A list containing a label vector, and a dgCMatrix object with 1611
Guolin Ke's avatar
Guolin Ke committed
96
97
98
99
#' rows and 126 variables
NULL

# Various imports
Guolin Ke's avatar
Guolin Ke committed
100
#' @import methods
Guolin Ke's avatar
Guolin Ke committed
101
102
#' @importFrom R6 R6Class
#' @useDynLib lightgbm
103
NULL