spark gbdt 自定义阈值 取出模型概率，并转换label

在大家使用spark 的时候，会发现，ML库下的模型生成的概率是一个Vector，那么如何将这个Vector的概率为1的那一列取出呢？并且自定义阈值，按照这个阈值切分得到label的0，1类呢？
这时候udf函数就派上了很好的用场，废话不多说，直接上代码

取出Vector 的第n列，生成新的dataframe

切分Vector得到每一列的值，形如调用下面的代码生成features,对feature进行切割

val vectorAssembler = new VectorAssembler()
  .setInputCols(int_col++dou_col).setOutputCol("features")

udf函数代码如下：

val code_par:(org.apache.spark.ml.linalg.Vector=>Double)=(arg:org.apache.spark.ml.linalg.Vector)=>{
      val temp =arg(1)
      temp
    }
    val parsecol=udf(code_par)
    val df_temp =df.withColumn(add_col,parsecol(df(“features”)))
    df_temp

当训练完成后，transform 预测数据后，也会生成一个形如上面这个features的parobability列，调用该udf函数，即可完成对label=1的prob的取出；

模型概率设置阈值，生成label

当得到label=1的prob时，使用udf函数，自定义设置阈值，多模型进行划分，并得到相应的精确率，recall；阈值划分的udf函数入下：

-- 概率阈值的切分
val thres = threshold
//阈值判断预测输出
val code :(Double => Int) = (arg: Double) => {if (arg < thres ) 0 else 1}
val colfun = udf(code)
val df_pre = finalDf.withColumn(predict_lable_col, col=colfun(finalDf(add_col)))

---

edited by :Eshter
date:20191012
版权归eshter 所有,转载请注明出处