Spark: How to get probabilities and AUC for Bernoulli Naive Bayes?

问题

I'm running a Bernoulli Naive Bayes using code:

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")

My question is how can I get the probability of membership to class 0 (or 1) and count AUC. I want to get similar result to LogisticRegressionWithSGD or SVMWithSGD where I was using this code:

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()

Unfortunately this code isn't working for NaiveBayes.

回答1:

Concerning the probabilities for Bernouilli Naive Bayes, here is an example :

// 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)}")

As for the 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()

Concerning the inquiry from the chat :

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)

and if you are only interested in the argmax classes :

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: Works with Spark 1.5+

EDIT: (for Pyspark users)

It seems like some are having troubles getting probabilities using pyspark and mllib. Well that's normal, spark-mllib doesn't present that function for pyspark.

Thus you'll need to use the spark-ml DataFrame-based 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       |
# +---------+-----+-----------------------------------------+----------------------------------------+----------+

You'll just need to select your prediction column and compute your AUC.

For more information about Naive Bayes in spark-ml, please refer to the official documentation here.

来源：https://stackoverflow.com/questions/33890062/spark-how-to-get-probabilities-and-auc-for-bernoulli-naive-bayes

标签

apache-spark

pyspark

apache-spark-mllib

naivebayes

apache-spark-ml