Merge pull request #109 from tidymodels/more-objectives-lightgbm

Author: EmilHvitfeldt

NEWS.md

@@ -1,5 +1,7 @@
 # bonsai (development version)
 
+* Automatic handling of `num_classes` argument when specifying a multiclass classification objective for the lightgbm engine (#109).
+
 * Increased the minimum R version to R 4.1.
 
 * Fixed bug where `num_threads` argument were ignored for lightgbm engine (#105).

R/lightgbm.R

@@ -101,7 +101,7 @@ train_lightgbm <- function(
     )
 
   args <- process_bagging(args)
-  args <- process_objective_function(args, x, y)
+  args <- process_objective_function(args, y)
 
   args <- sort_args(args)
 
@@ -175,24 +175,40 @@ process_mtry <- function(
   feature_fraction_bynode
 }
 
-process_objective_function <- function(args, x, y) {
+# https://lightgbm.readthedocs.io/en/latest/Parameters.html#core-parameters
+process_objective_function <- function(args, y) {
+  lvl <- levels(y)
+  lvls <- length(lvl)
   # set the "objective" param argument, clear it out from main args
   if (!any(names(args) %in% c("objective"))) {
     if (is.numeric(y)) {
       args$objective <- "regression"
     } else {
-      lvl <- levels(y)
-      lvls <- length(lvl)
       if (lvls == 2) {
-        args$num_class <- 1
         args$objective <- "binary"
       } else {
-        args$num_class <- lvls
         args$objective <- "multiclass"
       }
     }
   }
 
+  if (args$objective == "binary" && is.null(args$num_class)) {
+    args$num_class <- 1L
+  }
+
+  multiclass_obj <- c(
+    "multiclass",
+    "softmax",
+    "multiclassova",
+    "multiclass_ova",
+    "ova",
+    "ovr"
+  )
+
+  if (args$objective %in% multiclass_obj && is.null(args$num_class)) {
+    args$num_class <- lvls
+  }
+
   args
 }
 

tests/testthat/test-lightgbm.R

@@ -348,6 +348,30 @@ test_that("bonsai correctly determines objective when label is a factor", {
   expect_equal(bst$params$num_class, 3)
 })
 
+test_that("bonsai correctly determines num_classes when objective is set", {
+  skip_if_not_installed("lightgbm")
+  skip_if_not_installed("modeldata")
+
+  suppressPackageStartupMessages({
+    library(lightgbm)
+    library(dplyr)
+  })
+
+  data("penguins", package = "modeldata")
+  penguins <- penguins[complete.cases(penguins), ]
+
+  expect_no_error({
+    bst <- train_lightgbm(
+      x = penguins[, c("bill_length_mm", "bill_depth_mm")],
+      y = penguins[["species"]],
+      num_iterations = 5,
+      verbose = -1L,
+      objective = "multiclassova"
+    )
+  })
+  expect_identical(bst$params$objective, "multiclassova")
+  expect_identical(bst$params$num_class, 3L)
+})
 
 test_that("bonsai handles mtry vs mtry_prop gracefully", {
   skip_if_not_installed("modeldata")