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 所有,轉載請註明出處