Как GBTClassifier работает с несбалансированными данными для двоичной классификации?

Question

Я хотел выполнить двоичную классификацию с GBTClassifier для несбалансированного набора данных. Из документации spark я не видел ни одной опции, позволяющей это сделать.

Кто-нибудь имеет представление о том, как использовать GBTClassifier, указав тот факт, что наши данные не сбалансированы?

Спасибо

Примечание: я использую искру 2.3.2

Answer 1

Вот мое наивное решение: случайное сокращение выборки для большинства. Недостатком этого решения является потеря информации, и оно не будет работать для небольших наборов данных.

val resampledTrainDF = {

    val positiveLabel = "1"
    val trainDF_positives = trainDF.where(F.col(label) === positiveLabel)
    val trainDF_negatives = trainDF.where(F.col(label) =!= positiveLabel)

    val withReplacement = trainDF_positives.count >= trainDF_negatives.count

    if (withReplacement) {
        // downsampling positives
        val sampSize = math.round(  (1.0 * trainDF_negatives.count / trainDF_positives.count) * 1000) / 1000.0
        println("Downsampling Positives by " + (1 - sampSize)*100 + " %")
        trainDF_positives.sample(false, sampSize).union(trainDF_negatives)
    } else { 
        //downsampling negatives
        val sampSize = math.round(  (1.0 * trainDF_positives.count / trainDF_negatives.count) * 1000) / 1000.0
        println("Downsampling Negatives by " + (1 - sampSize)*100 +  "%")
        trainDF_negatives.sample(false, sampSize).union(trainDF_positives)
    }

}