Restore attributes in complete. (#5573)

This commit is contained in:
Jiaming Yuan 2020-04-22 02:06:55 +08:00 committed by GitHub
parent a734f52807
commit 564b22cee5
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 43 additions and 15 deletions

View File

@ -131,6 +131,25 @@ xgb.Booster.complete <- function(object, saveraw = TRUE) {
object$raw <- xgb.serialize(object$handle)
}
}
attrs <- xgb.attributes(object)
if (!is.null(attrs$best_ntreelimit)) {
object$best_ntreelimit <- as.integer(attrs$best_ntreelimit)
}
if (!is.null(attrs$best_iteration)) {
## Convert from 0 based back to 1 based.
object$best_iteration <- as.integer(attrs$best_iteration) + 1
}
if (!is.null(attrs$best_score)) {
object$best_score <- as.numeric(attrs$best_score)
}
if (!is.null(attrs$best_msg)) {
object$best_msg <- attrs$best_msg
}
if (!is.null(attrs$niter)) {
object$niter <- as.integer(attrs$niter)
}
return(object)
}

View File

@ -222,6 +222,15 @@ test_that("early stopping xgb.train works", {
early_stopping_rounds = 3, maximize = FALSE, verbose = 0)
)
expect_equal(bst$evaluation_log, bst0$evaluation_log)
xgb.save(bst, "model.bin")
loaded <- xgb.load("model.bin")
expect_false(is.null(loaded$best_iteration))
expect_equal(loaded$best_iteration, bst$best_ntreelimit)
expect_equal(loaded$best_ntreelimit, bst$best_ntreelimit)
file.remove("model.bin")
})
test_that("early stopping using a specific metric works", {