Нелинейная регрессия в Apache Spark. Разрабатываем своими руками09.12.2016 10:33

override protected def train(dataset: Dataset[_]): NonLinearRegressionModel = {
  // set instance weight to 1 if not defined the column
  val instanceWeightCol: Column = if (!isDefined(weightCol) || $(weightCol).isEmpty) lit(1.0) else col($(weightCol))
  val instances: RDD[Instance] = dataset.select(
    col($(labelCol)).cast(DoubleType), instanceWeightCol, col($(featuresCol))).rdd.map {
    case Row(label: Double, weight: Double, features: Vector) => Instance(label, weight, features)
  }

  // persists dataset if defined any storage level
  val handlePersistence = dataset.rdd.getStorageLevel == StorageLevel.NONE
  if (handlePersistence) instances.persist(StorageLevel.MEMORY_AND_DISK)

  val costFunc = new SquaresLossFunctionRdd(kernel, instances)
  val optimizer = new LBFGS[BDV[Double]]($(maxIter), 10, $(tol))

  // checks and assigns the initial coefficients
  val initial = {
    if (!isDefined(initCoeffs) || $(initCoeffs).length != kernel.dim) {
      if ($(initCoeffs).length != kernel.dim)
        logWarning(s"Provided initial coefficients vector (${$(initCoeffs)}) not corresponds with model dimensionality equals to ${kernel.dim}")
      BDV.zeros[Double](kernel.dim)
    } else
      BDV($(initCoeffs).clone())
  }

  val states = optimizer.iterations(new CachedDiffFunction[BDV[Double]](costFunc), initial)
  val (coefficients, objectiveHistory) = {
    val builder = mutable.ArrayBuilder.make[Double]
    var state: optimizer.State = null
    while (states.hasNext) {
      state = states.next
      builder += state.adjustedValue
    }

    // checks is method failed
    if (state == null) {
      val msg = s"${optimizer.getClass.getName} failed."
      logError(msg)
      throw new SparkException(msg)
    }

    (Vectors.dense(state.x.toArray.clone()).compressed, builder.result)
  }
  // unpersists the instances RDD
  if (handlePersistence) instances.unpersist()

  // produces the model and saves training summary
  val model = copyValues(new NonLinearRegressionModel(uid, kernel, coefficients))
  val (summaryModel: NonLinearRegressionModel, predictionColName: String) = model.findSummaryModelAndPredictionCol()
  val trainingSummary = new NonLinearRegressionSummary(
    summaryModel.transform(dataset), predictionColName, $(labelCol), $(featuresCol), objectiveHistory, model
  )
  model.setSummary(trainingSummary)
}

package org.apache.spark.ml.regression

import breeze.linalg.{DenseMatrix => BDM, DenseVector => BDV}

/**
  * Breeze implementation for the squares loss function.
  *
  * @param fitmodel concrete model implementation
  * @param xydata   labeled data combined into the matrix:
  *                 the first n-th columns consist of feature values, (n+1)-th columns contains labels
  */
class SquaresLossFunctionBreeze(val fitmodel: NonLinearFunction, xydata: BDM[Double])
  extends SquaresLossFunction {

  /**
    * The number of instances.
    */
  val instanceCount: Int = xydata.rows
  /**
    * The number of features.
    */
  val featureCount: Int = xydata.cols - 1
  /**
    * Feature matrix.
    */
  val X: BDM[Double] = xydata(::, 0 until featureCount)
  /**
    * Label vector.
    */
  val y: BDV[Double] = xydata(::, featureCount)

  /**
    * The model dimensionality (the number of weights).
    *
    * @return dimensionality
    */
  override def dim: Int = fitmodel.dim

  /**
    * Evaluates loss function value and the gradient vector
    *
    * @param weights weights
    * @return (loss function value, gradient vector)
    */
  override def calculate(weights: BDV[Double]): (Double, BDV[Double]) = {
    val r: BDV[Double] = diff(weights)
    (0.5 * (r.t * r), gradient(weights))
  }

  /**
    * Calculates a positive definite approximation of the Hessian matrix.
    *
    * @param weights weights
    * @return Hessian matrix approximation
    */
  override def hessian(weights: BDV[Double]): BDM[Double] = {
    val J: BDM[Double] = jacobian(weights)
    posDef(J.t * J)
  }

  /**
    * Calculates the Jacobian matrix
    *
    * @param weights weights
    * @return the Jacobian
    */
  def jacobian(weights: BDV[Double]): BDM[Double] = {
    val gradData = (0 until instanceCount) map { i => fitmodel.grad(weights, X(i, ::).t).toArray }
    BDM(gradData: _*)
  }

  /**
    * Calculates the difference vector between the label and the approximated values.
    *
    * @param weights weights
    * @return difference vector
    */
  def diff(weights: BDV[Double]): BDV[Double] = {
    val diff = (0 until instanceCount) map (i => fitmodel.eval(weights, X(i, ::).t) - y(i))
    BDV(diff.toArray)
  }

  /**
    * Calculates the gradient vector
    *
    * @param weights weights
    * @return gradient vector
    */
  def gradient(weights: BDV[Double]): BDV[Double] = {
    val J: BDM[Double] = jacobian(weights)
    val r = diff(weights)
    2.0 * J.t * r
  }
}

package org.apache.spark.ml.regression

import breeze.linalg.{DenseMatrix => BDM, DenseVector => BDV}
import org.apache.spark.broadcast.Broadcast
import org.apache.spark.ml.feature.Instance
import org.apache.spark.rdd.RDD

/**
  * Spark RDD implementation for the squares loss function.
  *
  * @param fitmodel concrete model implementation
  * @param xydata   RDD with instances
  */
class SquaresLossFunctionRdd(val fitmodel: NonLinearFunction, val xydata: RDD[Instance])
  extends SquaresLossFunction {

  /**
    * The model dimensionality (the number of weights).
    *
    * @return dimensionality
    */
  override def dim: Int = fitmodel.dim

  /**
    * Evaluates loss function value and the gradient vector
    *
    * @param weights weights
    * @return (loss function value, gradient vector)
    */
  override def calculate(weights: BDV[Double]): (Double, BDV[Double]) = {
    val bcW: Broadcast[BDV[Double]] = xydata.context.broadcast(weights)

    val (f: Double, grad: BDV[Double]) = xydata.treeAggregate((0.0, BDV.zeros[Double](dim)))(
      seqOp = (comb, item) => (comb, item) match {
        case ((loss, oldGrad), Instance(label, _, features)) =>
          val featuresBdv = features.asBreeze.toDenseVector
          val w: BDV[Double] = bcW.value
          val prediction = fitmodel.eval(w, featuresBdv)
          val addedLoss: Double = 0.5 * math.pow(label - prediction, 2)
          val addedGrad: BDV[Double] = 2.0 * (prediction - label) * fitmodel.grad(w, featuresBdv)
          (loss + addedLoss, oldGrad + addedGrad)
      },
      combOp = (comb1, comb2) => (comb1, comb2) match {
        case ((loss1, grad1: BDV[Double]), (loss2, grad2: BDV[Double])) => (loss1 + loss2, grad1 + grad2)
      })

    (f, grad)
  }

  /**
    * Calculates a positive definite approximation of the Hessian matrix.
    *
    * @param weights weights
    * @return Hessian matrix approximation
    */
  override def hessian(weights: BDV[Double]): BDM[Double] = {
    val bcW = xydata.context.broadcast(weights)

    val (hessian: BDM[Double]) = xydata.treeAggregate(new BDM[Double](dim, dim, Array.ofDim[Double](dim * dim)))(
      seqOp = (comb, item) => (comb, item) match {
        case ((oldHessian), Instance(_, _, features)) =>
          val grad = fitmodel.grad(bcW.value, features.asBreeze.toDenseVector)
          val subHessian: BDM[Double] = grad * grad.t
          oldHessian + subHessian
      },
      combOp = (comb1, comb2) => (comb1, comb2) match {
        case (hessian1, hessian2) => hessian1 + hessian2
      }
    )

    posDef(hessian)
  }
}