spark ml解决数据不均衡的简单方法

来源:
https://stackoverflow.com/questions/33372838/dealing-with-unbalanced-datasets-in-spark-mllib

数据不均衡有很多种解决方法,这里给一个比较简单的。
给dataframe增加一列权重列。
使用spark lr模型的时候设置参数setWeightCol填入权重列。

def balanceDataset(dataset: DataFrame): DataFrame = {

    // Re-balancing (weighting) of records to be used in the logistic loss objective function
    val numNegatives = dataset.filter(dataset("label") === 0).count
    val datasetSize = dataset.count
    val balancingRatio = (datasetSize - numNegatives).toDouble / datasetSize

    val calculateWeights = udf { d: Double =>
      if (d == 0.0) {
        1 * balancingRatio
      }
      else {
        (1 * (1.0 - balancingRatio))
      }
    }

    val weightedDataset = dataset.withColumn("classWeightCol", calculateWeights(dataset("label")))
    weightedDataset
  }

  val df_weighted = balanceDataset(df)
  val lr = new LogisticRegression().setLabelCol(labelCol).setWeightCol("classWeightCol")
©著作权归作者所有,转载或内容合作请联系作者
平台声明:文章内容(如有图片或视频亦包括在内)由作者上传并发布,文章内容仅代表作者本人观点,简书系信息发布平台,仅提供信息存储服务。

推荐阅读更多精彩内容