Spark：如何获得伯努利朴素贝叶斯的概率和 AUC？

我正在运行一个Bernoulli Naive Bayes使用代码：

val splits = MyData.randomSplit(Array(0.75, 0.25), seed = 2L)
val training = splits(0).cache()
val test = splits(1)
val model = NaiveBayes.train(training, lambda = 3.0, modelType = "bernoulli")

我的问题是如何获得 0（或 1）类成员资格的概率并计算 AUC。我想得到类似的结果LogisticRegressionWithSGD or SVMWithSGD我在哪里使用这段代码：

val numIterations = 100

val model = SVMWithSGD.train(training, numIterations)
model.clearThreshold()

// Compute raw scores on the test set.
val labelAndPreds = test.map { point =>
      val prediction = model.predict(point.features)
      (prediction, point.label)
}

// Get evaluation metrics.
val metrics = new BinaryClassificationMetrics(labelAndPreds)
val auROC = metrics.areaUnderROC() 

不幸的是这个代码不适用于NaiveBayes.

关于伯努利朴素贝叶斯的概率，这里有一个例子：

// Building dummy data
val data = sc.parallelize(List("0,1 0 0", "1,0 1 0", "1,0 0 1", "0,1 0 1","1,1 1 0"))

// Transforming dummy data into LabeledPoint
val parsedData = data.map { line =>
  val parts = line.split(',')
  LabeledPoint(parts(0).toDouble, Vectors.dense(parts(1).split(' ').map(_.toDouble)))
}

// Prepare data for training
val splits = parsedData.randomSplit(Array(0.75, 0.25), seed = 2L)
val training = splits(0).cache()
val test = splits(1)
val model = NaiveBayes.train(training, lambda = 3.0, modelType = "bernoulli")

// labels 
val labels = model.labels
// Probabilities for all feature vectors
val features = parsedData.map(lp => lp.features)
model.predictProbabilities(features).take(10) foreach println

// For one specific vector, I'm taking the first vector in the parsedData
val testVector = parsedData.first.features
println(s"For vector ${testVector} => probability : ${model.predictProbabilities(testVector)}")

至于AUC：

// Compute raw scores on the test set.
val labelAndPreds = test.map { point =>
  val prediction = model.predict(point.features)
  (prediction, point.label)
}

// Get evaluation metrics.
val metrics = new BinaryClassificationMetrics(labelAndPreds)
val auROC = metrics.areaUnderROC()

关于聊天中的询问：

val results = parsedData.map { lp =>
  val probs: Vector = model.predictProbabilities(lp.features)
  (for (i <- 0 to (probs.size - 1)) yield ((lp.label, labels(i), probs(i))))
}.flatMap(identity)

results.take(10).foreach(println)

// (0.0,0.0,0.59728640251696)
// (0.0,1.0,0.40271359748304003)
// (1.0,0.0,0.2546873180388961)
// (1.0,1.0,0.745312681961104)
// (1.0,0.0,0.47086939671877026)
// (1.0,1.0,0.5291306032812298)
// (0.0,0.0,0.6496075621805428)
// (0.0,1.0,0.3503924378194571)
// (1.0,0.0,0.4158585282373076)
// (1.0,1.0,0.5841414717626924)

如果您只对 argmax 类感兴趣：

val results = training.map { lp => val probs: Vector = model.predictProbabilities(lp.features)
  val bestClass = probs.argmax
  (labels(bestClass), probs(bestClass))
}
results.take(10) foreach println

// (0.0,0.59728640251696)
// (1.0,0.745312681961104)
// (1.0,0.5291306032812298)
// (0.0,0.6496075621805428)
// (1.0,0.5841414717626924)

Note:与火花1.5+

编辑：（对于 Pyspark 用户）

似乎有些人在使用概率获取时遇到了麻烦pyspark and mllib。嗯，这很正常，Spark-MLlib不为 pyspark 提供该功能。

因此你需要使用spark-ml DataFrame基于API：

from pyspark.sql import Row
from pyspark.ml.linalg import Vectors
from pyspark.ml.classification import NaiveBayes

df = spark.createDataFrame([
    Row(label=0.0, features=Vectors.dense([0.0, 0.0])),
    Row(label=0.0, features=Vectors.dense([0.0, 1.0])),
    Row(label=1.0, features=Vectors.dense([1.0, 0.0]))])

nb = NaiveBayes(smoothing=1.0, modelType="bernoulli")
model = nb.fit(df)

model.transform(df).show(truncate=False)
# +---------+-----+-----------------------------------------+----------------------------------------+----------+
# |features |label|rawPrediction                            |probability                             |prediction|
# +---------+-----+-----------------------------------------+----------------------------------------+----------+
# |[0.0,0.0]|0.0  |[-1.4916548767777167,-2.420368128650429] |[0.7168141592920354,0.28318584070796465]|0.0       |
# |[0.0,1.0]|0.0  |[-1.4916548767777167,-3.1135153092103742]|[0.8350515463917526,0.16494845360824742]|0.0       |
# |[1.0,0.0]|1.0  |[-2.5902671654458262,-1.7272209480904837]|[0.29670329670329676,0.7032967032967034]|1.0       |
# +---------+-----+-----------------------------------------+----------------------------------------+----------+

您只需要选择您的预言列并计算您的 AUC。

关于spark-ml中朴素贝叶斯的更多信息，请参考官方文档here.