深入理解Spark ML：多项式朴素贝叶斯原理与源码分析

import org.apache.spark.ml.classification.NaiveBayes

// 加载数据
val data = spark.read.format("libsvm").load("data/mllib/sample_libsvm_data.txt")

// 切分数据集与训练集
val Array(trainingData, testData) = data.randomSplit(Array(0.7, 0.3), seed = 1234L)

// 训练朴素贝叶斯模型
val model = new NaiveBayes()
  .fit(trainingData)

// 预测
val predictions = model.transform(testData)
predictions.show()

  override protected def train(dataset: Dataset[_]): NaiveBayesModel = {
    trainWithLabelCheck(dataset, positiveLabel = true)
  }

 private[spark] def trainWithLabelCheck(
      dataset: Dataset[_],
      positiveLabel: Boolean): NaiveBayesModel = {
      //检测label
    if (positiveLabel && isDefined(thresholds)) {
      val numClasses = getNumClasses(dataset)
      require($(thresholds).length == numClasses, this.getClass.getSimpleName +
        ".train() called with non-matching numClasses and thresholds.length." +
        s" numClasses=$numClasses, but thresholds has length ${$(thresholds).length}")
    }
    //模型类型 多项式朴素贝叶斯是  Multinomial
    val modelTypeValue = $(modelType)
    val requireValues: Vector => Unit = {
      modelTypeValue match {
        case Multinomial =>
          // 确认所有的值非负
          // values.forall(_ >= 0.0)
          requireNonnegativeValues
        ......
      }
    }
    // Instrumentation 是 一个小封装，用来定义为一个estimator定义一个training session和该session中有学用的信息的log方法
    val instr = Instrumentation.create(this, dataset)
    instr.logParams(labelCol, featuresCol, weightCol, predictionCol, rawPredictionCol,
      probabilityCol, modelType, smoothing, thresholds)
    // 得到特征维度数，即公式中的 n
    val numFeatures = dataset.select(col($(featuresCol))).head().getAs[Vector](0).size
    instr.logNumFeatures(numFeatures)
    // 得到记录的权重 为设置 默认为 1.0
    val w = if (!isDefined(weightCol) || $(weightCol).isEmpty) lit(1.0) else col($(weightCol))

    // 聚合
    val aggregated = dataset.select(col($(labelCol)), w, col($(featuresCol))).rdd
      .map { row => (row.getDouble(0), (row.getDouble(1), row.getAs[Vector](2)))
      // 根据key labelCol 进行聚合
      // value 的初始值为 0.0,Vectors.zeros(numFeatures).toDense
      }.aggregateByKey[(Double, DenseVector)]((0.0, Vectors.zeros(numFeatures).toDense))(
      // 合并在同一个partition中的值
      seqOp = {
         case ((weightSum: Double, featureSum: DenseVector), (weight, features)) =>
           requireValues(features)
           // featureSum = featureSum + weight × features
           BLAS.axpy(weight, features, featureSum)
           (weightSum + weight, featureSum)
      },
      //合并不同partition中的值
      combOp = {
         case ((weightSum1, featureSum1), (weightSum2, featureSum2)) =>
           BLAS.axpy(1.0, featureSum2, featureSum1)
           (weightSum1 + weightSum2, featureSum1)
      }).collect().sortBy(_._1)

    // label 的类别数，即公式中的 K
    val numLabels = aggregated.length
    instr.logNumClasses(numLabels)
    // 文档数，即公式中的 N
    val numDocuments = aggregated.map(_._2._1).sum

    val labelArray = new Array[Double](numLabels)
    // 用来存 log(p(y_i)
    val piArray = new Array[Double](numLabels)
    // 用来存 log(P(a_{j}|y_i))
    val thetaArray = new Array[Double](numLabels * numFeatures)

    val lambda = $(smoothing)
    // log(N+Klambda)
    val piLogDenom = math.log(numDocuments + numLabels * lambda)
    var i = 0
    aggregated.foreach { case (label, (n, sumTermFreqs)) =>
      labelArray(i) = label
      // log(p(y_i)) = log(\sum^N_{i=1}I(y_i=c_k)+lambda) - log(N+Klambda)
      piArray(i) = math.log(n + lambda) - piLogDenom
      // log(\sum_{i=1}^N\sum_{j=1}^na_{j,y_i=c_k}+nlambda)
      val thetaLogDenom = $(modelType) match {
        case Multinomial => math.log(sumTermFreqs.values.sum + numFeatures * lambda)
       ......
      }
      var j = 0
      while (j < numFeatures) {
      // log(P(a_{j}|y_i))=log(\sum_{i=1}^Na_{j,y_i=c_k}+lambda)-log(\sum_{i=1}^N\sum_{j=1}^na_{j,y_i=c_k}+nlambda)
        thetaArray(i * numFeatures + j) = math.log(sumTermFreqs(j) + lambda) - thetaLogDenom
        j += 1
      }
      i += 1
    }

    val pi = Vectors.dense(piArray)
    val theta = new DenseMatrix(numLabels, numFeatures, thetaArray, true)
    // 生成并返回模型
    val model = new NaiveBayesModel(uid, pi, theta).setOldLabels(labelArray)
    instr.logSuccess(model)
    model
  }

  override protected def predict(features: FeaturesType): Double = {
    raw2prediction(predictRaw(features))
  }

// 返回之大值的坐标
protected def raw2prediction(rawPrediction: Vector): 
Double = rawPrediction.argmax

  override protected def predictRaw(features: Vector): Vector = {
    $(modelType) match {
      case Multinomial =>
        multinomialCalculation(features)
      ......
    }
  }

  private def multinomialCalculation(features: Vector) = {
    // 求各个label类别下的条件概率
    val prob = theta.multiply(features)
    // 加上先验概率：prob = prob + 1.0 * pi
    BLAS.axpy(1.0, pi, prob)
    prob
  }