如何只得到概率大于 x 的预测

Question

我使用随机森林将文本分类到特定类别。当我使用我的测试数据时，我得到了 0.98 的准确度。但是对于另一组数据，整体准确度下降到 0.7。我认为，大多数行仍然具有很高的准确性。

所以现在我只想显示具有高置信度的预测类别。 random-forrest 给了我一列 "probability"，这是一个概率数组。如何获得所选预测的实际概率？

val randomForrest = new RandomForestClassifier()
      .setLabelCol(labelIndexer.getOutputCol)
      .setFeaturesCol(vectorAssembler.getOutputCol)
      .setProbabilityCol("probability")
      .setSeed(123)
      .setPredictionCol("prediction")

Answer 1

我最终想出了以下 udf 来获得最佳预测及其概率。如果有更方便的方法欢迎评论

def getBestPrediction = udf((
  rawPrediction: org.apache.spark.ml.linalg.Vector, probability: org.apache.spark.ml.linalg.Vector) => {
  val bestPrediction = probability.argmax
  val bestProbability = probability(bestPrediction)     
  (bestPrediction, bestProbability)
})

如何只得到概率大于 x 的预测

How to get only the predictions with a probability greater than x

random-forest

apache-spark-mllib