allow watchlist to be NamedTuple (#199)

dmlc · Nov 21, 2023 · 17f303a · 17f303a
1 parent 0b4977a
commit 17f303a
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 5 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,6 +1,6 @@
 name = "XGBoost"
 uuid = "009559a3-9522-5dbb-924b-0b6ed2b22bb9"
-version = "2.5.0"
+version = "2.5.1"
 
 [deps]
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"

diff --git a/src/XGBoost.jl b/src/XGBoost.jl
@@ -23,7 +23,6 @@ include("Lib.jl")
 using .Lib
 using .Lib: DMatrixHandle, BoosterHandle
 
-
 const LOG_LEVEL_REGEX = r"\[.*\] (\D*): "
 
 function xgblog(s::Cstring)

diff --git a/src/booster.jl b/src/booster.jl
@@ -428,7 +428,7 @@ for custom loss.
 """
 function update!(b::Booster, data, a...;
                  num_round::Integer=1, 
-                 watchlist::Any = Dict("train" => data), 
+                 watchlist=Dict("train" => data), 
                  early_stopping_rounds::Integer=0,
                  maximize=false,
                  kw...,
@@ -578,7 +578,7 @@ ŷ = predict(b, dvalid, ntree_limit = b.best_iteration)
 """
 function xgboost(dm::DMatrix, a...;
                 num_round::Integer=10,
-                watchlist::AbstractDict = Dict("train" => dm),
+                watchlist=Dict("train" => dm),
                 early_stopping_rounds::Integer=0, 
                 maximize=false,
                 kw...
@@ -590,7 +590,7 @@ function xgboost(dm::DMatrix, a...;
     # We have a watchlist - give a warning if early stopping is provided and watchlist is a Dict type with length > 1
     if isa(watchlist, Dict)
         if early_stopping_rounds > 0 && length(watchlist) > 1
-            error("You must supply an OrderedDict type for watchlist if early stopping rounds is enabled and there is more than one element in watchlist.")
+            error("You must supply an OrderedDict or NamedTuple type for watchlist if early stopping rounds is enabled and there is more than one element in watchlist.")
         end
     end
 

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -183,7 +183,16 @@ end
         early_stopping_rounds = 2
         )
 
+    watchlist_nt = (train=dtrain, eval=dtest)
 
+    bst_early_stopping = xgboost(dtrain,
+        num_round=30,
+        watchlist=watchlist_nt,
+        η=1,
+        objective="binary:logistic",
+        eval_metric=["rmsle","rmse"],
+        early_stopping_rounds = 2
+        )
 
     @test XGBoost.getnrounds(bst_early_stopping) > 2
     @test XGBoost.getnrounds(bst_early_stopping) <= nrounds_bst