Commit 95f213bd authored by Laurae's avatar Laurae Committed by Guolin Ke
Browse files

Add binary/classification requirement for strata (#897)

parent de51990e
...@@ -323,7 +323,7 @@ generate.cv.folds <- function(nfold, nrows, stratified, label, group, params) { ...@@ -323,7 +323,7 @@ generate.cv.folds <- function(nfold, nrows, stratified, label, group, params) {
rnd_idx <- sample(seq_len(nrows)) rnd_idx <- sample(seq_len(nrows))
# Request stratified folds # Request stratified folds
if (isTRUE(stratified) && length(label) == length(rnd_idx)) { if (isTRUE(stratified) && params$objective %in% c("binary", "multiclass") && length(label) == length(rnd_idx)) {
y <- label[rnd_idx] y <- label[rnd_idx]
y <- factor(y) y <- factor(y)
...@@ -400,6 +400,7 @@ lgb.stratified.folds <- function(y, k = 10) { ...@@ -400,6 +400,7 @@ lgb.stratified.folds <- function(y, k = 10) {
} }
if (k < length(y)) { if (k < length(y)) {
## Reset levels so that the possible levels and ## Reset levels so that the possible levels and
## the levels in the vector are the same ## the levels in the vector are the same
y <- factor(as.character(y)) y <- factor(as.character(y))
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment