如何分割ML桩线Logistic回归产生的预测概率

时间:2015-09-14 17:48:10

标签: scala apache-spark apache-spark-ml

我试图使用ML管道和DataFrame API从逻辑模型中提取预测概率。预测概率的输出是列向量,其存储每个类(0,1)的预测概率,如下所示。我想知道如何才能提取第1类的概率。谢谢!

  

概率
  " [0.13293408418007766,0.8670659158199223]"
  " [0.1335112097146626,0.8664887902853374]"

3 个答案:

答案 0 :(得分:5)

像这样的UDF应该可以工作:

import org.apache.spark.sql.functions.udf

val getPOne = udf((v: org.apache.spark.mllib.linalg.Vector) => v(1))
model.transform(testDf).select(getPOne($"probability"))

答案 1 :(得分:0)

这样的方式也可以达到目标。

import org.apache.spark.mllib.linalg.Vector
model.transform(dataDF).select("prob").map(row => row.getAs[Vector]("prob").toArray)

那么你现在可以得到你想要的东西。

答案 2 :(得分:0)

您还可以先将概率列添加到org.apache.spark.mllib.linalg.DenseVector,然后提取该元素。

model.transform(testDf).select("probability").map(e=> e.asInstanceOf[DenseVector]).map(e=> e(1))