如何拆分 ML 管道逻辑回归产生的预测概率

How to Split the Predicted Probabilities Produced by ML Pileline Logistic Regression

我正在尝试使用 ML 管道和 DataFrame API 从逻辑模型中提取预测概率。预测概率的输出是一个列向量,它存储每个 class(0, 1) 的预测概率,如下所示。我想知道如何只提取 class 1 的概率。谢谢!

prob
"[0.13293408418007766,0.8670659158199223]"
"[0.1335112097146626,0.8664887902853374]"

像这样的 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"))

这样也能达到目的

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

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

也可以先将概率列大小写为org.apache.spark.mllib.linalg.DenseVector,再提取元素。

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