Spark ML 朴素贝叶斯用概率预测多个类别

Question

有没有办法让模型返回预测标签列表以及每个标签的概率得分？

例如 给定特征（f1，f2，f3）， 它返回类似这样的内容： 标签1：0.50，标签2：0.33...

在 Spark 中可行吗？

Answer 1

是的，这是可能的。

rawPrediction

列的输出是

Array[Double]

，其中包含每个标签的概率。

在您的示例中，此列将是一个数组（0.5,0.33,0.17），您必须编写一个 UDF 将此数组转换为字符串。

需要注意的是，如果您使用 StringIndexer 对标签列进行编码，则生成的标签将与原始标签不同。 （最常见的标签获得索引 0）

有一些代码可以执行类似的操作，可以适应您的用例。 我的代码只是将每个功能的前 X 个预测写入 CSV 文件。 writeToCsv 的参数 @df 必须是经过朴素贝叶斯模型转换后的 DataFrame。

 def topXPredictions(v: Vector, labels: Broadcast[Array[String]], topX: Int): Array[String] = {
    val labelVal = labels.value
    v.toArray
      .zip(labelVal)
      .sortBy {
        case (score, label) => score
      }
      .reverse
      .map {
        case (score, label) => label
      }
      .take(topX)
  }

  def writeToCsv(df: DataFrame, labelsBroadcast: Broadcast[Array[String]], name: String = "output"): Unit = {
    val get_top_predictions = udf((v: Vector, x: Int) => topXPredictions(v, labelsBroadcast, x))

      df
      .select(
        col("id")
        ,concat_ws(" ", get_top_predictions(col("rawPrediction"), lit(10))).alias("top10Predictions")
       )
      .orderBy("id")
      .coalesce(1)
      .write
      .mode(SaveMode.Overwrite)
      .format("com.databricks.spark.csv")
      .option("header", "true")
      .save(name)
  }

Answer 2

虽然这个问题很旧并且已经得到解答，但我将在 2024 年执行以下操作（使用 Spark 3.5.1 中的

org.apache.spark.ml.classification.NaiveBayes

）：

nb = NaiveBayes()
model = nb.fit(features)
predictionsTraining = model.transform(features)

predictionsTraining

包含

features

中的所有列以及以下列：

• rawPrediction

  ：一组 NaiveBayes 值（每个类一个值）

• probability

  ：一组概率（每个类一个）；概率之和为 1。

• prediction

  ：预测类别（即

  probability

  具有最高值的类别）

注意：为了简单起见，我们预测与训练相同的特征（您不应该在现实世界中这样做）。