Spark ML库简析

发表于 |

讨论Spark上分布式机器学习库的实现。
Spark的机器学习库主要分为ml和mllib，其中ml较新，本文主要围绕ml来讲。
ml和mllib底层会用到Breeze库（类似于numpy的线性代数库）和BLAS（Basic Linear Algebra Subroutines，更基础的线性代数库），在这里也进行介绍。

LR回归的实现

总体实现

ml库的整个训练和预测过程可以用下面三行代码来概括

1
2
3
val classifier = new LogisticRegression()
val model：LogisticRegressionModel = classifier.fit(train)
val predictions = model.transform(test)

训练过程

LogisticRegression类的训练入口是fit。fit首先定义在基类Predictor.scala中，它会首先做一些事情，比如转换labelCol的类型到Double，然后调用train。train会返回一个Model，但在Predictor中并没有实现，而是交付给子类来实现。
LogisticRegression类的train是自己实现的。主要包括下面几个部分。

  1. 校验打印参数
    参数包含下面的一些列名

    1
    labelCol, weightCol, featuresCol, predictionCol, rawPredictionCol, probabilityCol

    这里的weightCol可以给样本指定权重，用来解决数据不平衡问题。详见balanceDataset的实现。
    还包括下面的一些训练用参数

    1
    regParam, elasticNetParam, standardization, threshold, maxIter, tol, fitIntercept

  2. 计算summarizer和labelSummarizer
    可以得到一些基础信息，比如直方图、特征数量、特征的均值方差和数量，class数量。其中直方图(histogram)，大概内容类似[0: 34.0, 1: 66.0]，意味着分类1拥有34样本，分类2有66个样本。均值(mean)大概内容类似[0: -0.03327765069036007 ]。
    根据这些信息还会判断任务类型，是不是多类的（设置isMultinomial）

  3. 计算coefficientMatrix、interceptVector和objectiveHistory
    这是主要过程，分别是返回系数矩阵、偏置值向量和loss之类的东西。
    返回值是(denseCoefficientMatrix.compressed,interceptVec.compressed,arrayBuilder.result())

  4. 计算summaryModel、probabilityColName、predictionColName

模型表示

模型存储

首先Model会继承一个org.apache.spark.ml.util.MLWritable。这样就支持model.save()方法。另外，可以通过MLWriter来保存，对应model.write.overwrite().save()

预测过程

LogisticRegressionModel的transform方法实际上实现在ProbabilisticClassificationModel里面，最终会调用predictRaw计算rawPrediction列。predictRaw和下面提到的一些方法都定义在ClassificationModel里面，不过是虚的。

1
2
3
4
5
6
7
if ($(rawPredictionCol).nonEmpty) {
  val predictRawUDF = udf { (features: Any) =>
    predictRaw(features.asInstanceOf[FeaturesType])
  }
  outputData = outputData.withColumn(getRawPredictionCol, predictRawUDF(col(getFeaturesCol)))
  numColsOutput += 1
}

然后再调用predictProbability或者raw2probability计算probability列，这个列是一个Vector，表示每一个label的probability，加起来是等于1的（见ProbabilisticClassificationModel里面的注释）。所有继承了ProbabilisticClassifier的分类器都会有在这一列。

1
2
3
4
5
6
7
8
9
10
11
12
if ($(probabilityCol).nonEmpty) {
  val probUDF = if ($(rawPredictionCol).nonEmpty) {
    udf(raw2probability _).apply(col($(rawPredictionCol)))
  } else {
    val probabilityUDF = udf { (features: Any) =>
      predictProbability(features.asInstanceOf[FeaturesType])
    }
    probabilityUDF(col($(featuresCol)))
  }
  outputData = outputData.withColumn($(probabilityCol), probUDF)
  numColsOutput += 1
}

然后再调用raw2prediction或者probability2prediction或者predict计算prediction列。

训练过程详解

Summary过程

首先使用MetadataUtils.getNumClasses来获得Class数量。

1
2
3
4
5
6
7
val numClasses = MetadataUtils.getNumClasses(dataset.schema($(labelCol))) match {
  case Some(n: Int) =>
    require(n >= histogram.length, s"Specified number of classes $n was " +
      s"less than the number of unique labels ${histogram.length}.")
    n
  case None => histogram.length
}

MetadataUtils是一个私有对象，里面会计算得到labelSchema对应的Attribute，然后来判断有多少类。Attribute对象的实现比较复杂，单独提到后面讲。

1
2
3
4
5
6
7
def getNumClasses(labelSchema: StructField): Option[Int] = {
  Attribute.fromStructField(labelSchema) match {
    case binAttr: BinaryAttribute => Some(2)
    case nomAttr: NominalAttribute => nomAttr.getNumValues
    case _: NumericAttribute | UnresolvedAttribute => None
  }
}

接着判断是否是多类分类任务，主要是根据family的值来判断，再根据numClasses的值来校验。family可以取下面的值:

  1. auto（默认值）
    自动根据class数量选择，如果numClasses == 1 || numClasses == 2，则是binomial，否则是multinomial
  2. binomial
    基于pivoting的LR
  3. multinomial
    softmax的LR，不基于pivoting
    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    @Since("2.1.0")
    final val family: Param[String] = new Param(this, "family",
      "The name of family which is a description of the label distribution to be used in the " +
        s"model. Supported options: ${supportedFamilyNames.mkString(", ")}.",
      (value: String) => supportedFamilyNames.contains(value.toLowerCase(Locale.ROOT)))

    /** @group getParam */
    @Since("2.1.0")
    def getFamily: String = $(family)

    val isMultinomial = getFamily.toLowerCase(Locale.ROOT) match {
      case "binomial" =>
        require(numClasses == 1 || numClasses == 2, s"...")
        false
      case "multinomial" => true
      case "auto" => numClasses > 2
      case other => throw new IllegalArgumentException(s"Unsupported family: $other")
    }
    val numCoefficientSets = if (isMultinomial) numClasses else 1

优化过程

主要是创建了损失函数costFun，并对于这个损失函数使用优化器进行优化。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
val regParamL1 = $(elasticNetParam) * $(regParam)
val regParamL2 = (1.0 - $(elasticNetParam)) * $(regParam)
val getAggregatorFunc = new LogisticAggregator(bcFeaturesStd, numClasses, $(fitIntercept),
  multinomial = isMultinomial)(_)
val regularization = if (regParamL2 != 0.0) {
  val shouldApply = (idx: Int) => idx >= 0 && idx < numFeatures * numCoefficientSets
  Some(new L2Regularization(regParamL2, shouldApply,
    if ($(standardization)) None else Some(getFeaturesStd)))
} else {
  None
}
val costFun = new RDDLossFunction(instances, getAggregatorFunc, regularization,
  $(aggregationDepth))
val states = optimizer.iterations(new CachedDiffFunction(costFun),
  new BDV[Double](initialCoefWithInterceptMatrix.toArray))

首先来看损失函数RDDLossFunction的类定义。
首先ClassTag用来实现保障类型擦除后类型安全的功能。

1
2
3
4
5
6
7
8
private[ml] class RDDLossFunction[
    T: ClassTag,
    Agg <: DifferentiableLossAggregator[T, Agg]: ClassTag](
    instances: RDD[T],
    getAggregator: (Broadcast[Vector] => Agg),
    regularization: Option[DifferentiableRegularization[Vector]],
    aggregationDepth: Int = 2)
  extends DiffFunction[BDV[Double]] {
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
  override def calculate(coefficients: BDV[Double]): (Double, BDV[Double]) = {
    val bcCoefficients = instances.context.broadcast(Vectors.fromBreeze(coefficients))
    val thisAgg = getAggregator(bcCoefficients)
    val seqOp = (agg: Agg, x: T) => agg.add(x)
    val combOp = (agg1: Agg, agg2: Agg) => agg1.merge(agg2)
    val newAgg = instances.treeAggregate(thisAgg)(seqOp, combOp, aggregationDepth)
    val gradient = newAgg.gradient
    val regLoss = regularization.map { regFun =>
      val (regLoss, regGradient) = regFun.calculate(Vectors.fromBreeze(coefficients))
      BLAS.axpy(1.0, regGradient, gradient)
      regLoss
    }.getOrElse(0.0)
    bcCoefficients.destroy()
    (newAgg.loss + regLoss, gradient.asBreeze.toDenseVector)
  }
}

优化器

LBGFS

牛顿法的特点是收敛很快，但是运用牛顿法需要计算二阶偏导数，而且目标函数的Hesse矩阵可能非正定。对此，可以使用拟牛顿法，也就是用不包含二阶导数的矩阵近似牛顿法中的Hesse矩阵的逆矩阵。 是一种一阶方法。由于构造近似矩阵的方法不同，因而出现不同的拟牛顿法，BFGS是其中一种。因为BFGS仍然存在较大的内存占用。因此有了LBFGS算法，其中L指的是Limited的意思。
下面，简单查看一下LBFGS的代码，可以发现，它继承了FirstOrderMinimizer优化器，印证了LBFGS是一个一阶方法。

1
2
3
class LBFGS[T](convergenceCheck: ConvergenceCheck[T], m: Int)(implicit space: MutableInnerProductModule[T, Double])
    extends FirstOrderMinimizer[T, DiffFunction[T]](convergenceCheck)
    with SerializableLogging {

OWLQN

OWL-QN(Orthant-Wise Limited-Memory Quasi-Newton)算法，该算法是基于L-BFGS算法的可用于求解L1正则的算法。

Loss函数

treeAggregate

seqOp操作会聚合各分区中的元素，然后combOp操作把所有分区的聚合结果再次聚合，两个操作的初始值都是zeroValue。treeAggregate和aggregate的区别是，aggregate在每个分区处理完之后就直接交给Driver合并了，但treeAggregate会在Executor直接树形地Aggregate。

1
2
3
4
5
6
7
8
9
10
val sc = spark.sparkContext
val rdd = sc.parallelize(1 to 100000).repartition(6)
val start = System.currentTimeMillis();
val reduceResult = rdd.reduce{(x, y) => x + y}
val endReduce = System.currentTimeMillis();
val aggregateResult = rdd.aggregate(0)((x, y) => x + y, (x, y) => x + y)
val endAgg = System.currentTimeMillis();
val treeAggregateResult = rdd.treeAggregate(0)((x, y) => x + y, (x, y) => x + y)
val endTreeAgg = System.currentTimeMillis();
println(s"reduceResult ${reduceResult} ${(endReduce-start).toString} aggregateResult ${aggregateResult} ${(endAgg-endReduce).toString} treeAggregateResult ${treeAggregateResult} ${(endTreeAgg-endAgg).toString}")

在六台机器上跑了，耗时如下

1
reduceResult 705082704 7073 aggregateResult 705082704 831 treeAggregateResult 705082704 755

BLAS库

BLAS库的选择

Spark的BLAS库实现在mllib/linalg下面。其底层借助于com.github.fommil.netlib，而netelib有两个BLAS库的实现，其中f2jBLAS是从Fortran代码翻译到Java里面的，nativeBLAS是用系统原生的。两个库的使用规则如下所示，至于为什么，参考爆栈网的解释的意思是f2j的实现在Level较低的情况的性能要更好一点。

1
2
3
4
5
6
7
private[mllib] def getBLAS(vectorSize: Int): NetlibBLAS = {
  if (vectorSize < nativeL1Threshold) {
    f2jBLAS
  } else {
    nativeBLAS
  }
}

这里的Level，指的是BLAS函数分为三个Level，Level越高，计算越复杂。Level1是向量之间的操作，Level2是 向量和矩阵之间的操作，Level3是矩阵和矩阵之间的操作。

axpy和Vector的实现

axpy指的是线代里面的$y = a * x$这样的操作。
axpy里面根据是否稀疏，进行讨论。对于DenseVector会直接调用daxpy。这里daxpy的d不是dense，而是指double，对应的saxpy指的是单精度，还有caxpy代表复数单精度和zaxpy代表复数双精度。
对于SparseVector则要首先进行解码为DenseVector。SparseVector 用两个数组，一个记录原始向量中非零元素的值，另一个记录非零元素对应到原始向量的位置。这有点类似于一维的COO表示形式。

但不是所有的计算都需要转换为DenseVector的，例如向量的dot就不需要。观摩下面的代码发现，稀疏向量x中为0的列就不需要乘了。事实上有人进行了比较，Sparse算得快，但是构造慢。其中nnz表示number of nonzero entries。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
private def dot(x: SparseVector, y: DenseVector): Double = {
  val xValues = x.values
  val xIndices = x.indices
  val yValues = y.values
  val nnz = xIndices.length

  var sum = 0.0
  var k = 0
  while (k < nnz) {
    sum += xValues(k) * yValues(xIndices(k))
    k += 1
  }
  sum
}

gemm和Matrix的实现

gemm，即通用矩阵乘（GEMM，General Matrix Multiplication），其公式为$ C = alpha * A * B + beta * C $，其定义如下。其中A大小$m*k$，B大小$k*n$，C大小$m*n$，且C.isTransposed必须是false。根据A是否为Dense，可以分为2个版本。

1
2
3
def gemm(alpha: Double,
    A: Matrix, B: DenseMatrix,
    beta: Double, C: DenseMatrix): Unit = {

这里的Matrix同样分为Dense和Sparse。isTransposed表示有没有进行过转置。我理解用这个的原因是为了实现lazy的转置。
DenseMatrix的四要素是numRows、numCols、 values和isTransposed，values是一维数组，因为我们可以通过numRows和numCols找到定位。
SparseMatrix是六要素，多了colPtrs和rowIndices，values是一维数组，按列存储，实际上就是Compressed Sparse Columns(CSC)格式存储，如果isTransposed是true，那么就按照Compressed Sparse Row(CSR)存储。以isTransposed为false的情况为例，colPtrs表示每一列的起始在values中的位置，rowIndices表示在values中的每一个元素在是这一列中的第几行。

下面展示一下CSC的表示方法，例如矩阵

1
2
3
4
5
6
1.0 0.0 4.0
0.0 3.0 5.0
2.0 0.0 6.0
values=[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
rowIndices=[0, 2, 1, 0, 1, 2]
colPointers=[0, 2, 3, 6]

下面看gemm的实现，首先要对A是不是Dense进行讨论，对Dense而言是平凡的，调用dgemm就行了。对于Spark而言，需要根据A和B有没有transpose过的四种情况进行讨论。下面我们来看A是Sparse，B是Dense对应的gemm重载下的四种情况。

首先来看A和B都没有转置的情况，我们现在是需要用A的对应行乘以B的对应列，可以看做左边的矩阵的每一列，分别去乘上右边矩阵对应的行，得到n个矩阵，然后我们将它们对应位置加起来就能得到C。这种方式可以参考我的MIT线性代数的学习报告

首先先要放缩一下C，这样我们的列乘行的操作可以直接在C.values上面累加，节约空间。

1
2
3
if (beta != 1.0) {
  getBLAS(C.values.length).dscal(C.values.length, beta, C.values, 1)
}

下面是主体，包含两个循环。外层循环变量colCounterForB每轮递增，用来枚举到B的所有列nB，内层循环变量colCounterForA用来枚举A的所有列

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
while (colCounterForB < nB) {
  var colCounterForA = 0 // The column of A to multiply with the row of B
  // kB=B行数=A列数kA
  // 因为按列存储，BStart表示第colCounterForB列开始的位置。
  // 因为B是Dense的，所以可以直接乘
  val Bstart = colCounterForB * kB
  // mA = A行数
  // Cstart表示第colCounterForB列开始的位置，是对C而言的，且C有mA行
  val Cstart = colCounterForB * mA
  
  while (colCounterForA < kA) {
    // 计算C的第colCounterForB列上所有行的部分值
    var i = AcolPtrs(colCounterForA)
    val indEnd = AcolPtrs(colCounterForA + 1)
    // indEnd-i表示第colCounterForA列有多少个元素
    // Avals(i)于对CSC格式的A的行进行遍历
    // BVal是B的第colCounterForB列第colCounterForA行
    // 遍历A的colCounterForA列上的每个行元素
    // 可以看到，Cstart在内层循环是不变的，所以对Cstart对应的第colCounterForB列的操作会重复kA次
    val Bval = Bvals(Bstart + colCounterForA) * alpha
    while (i < indEnd) {
      // Cvals是c.values
      // 对于A的colCounterForA上的每个行元素，加到C的对应位置上面，注意C是Dense的。
      // 这Avals(i)相当于对CSC格式的行进行遍历
      Cvals(Cstart + ArowIndices(i)) += Avals(i) * Bval
      i += 1
    }
    colCounterForA += 1
  }
  colCounterForB += 1
}

现在来看A转置B没有转置的情况，。我们现在是需要用A的对应行乘以B的对应列。因为A被转置了，所以问题实际上变简单了，因为A转置后的行，实际上就是A转置前的列，在物理上是连续存储的。因为我们恰恰需要的是A转置后的行，所以实际可以当做转置前的列来处理。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
// colCounterForB在每个while后递增，枚举到B的所有列nB
while (colCounterForB < nB) {
  var rowCounterForA = 0
  // mA=A行数，因为按列存储，表示第colCounterForB列开始的位置
  val Cstart = colCounterForB * mA 
  // kA=A列数=B行数kB，实际上是B的第colCounterForB列开始的位置
  val Bstart = colCounterForB * kA 
  // rowCounterForA枚举A的所有行
  while (rowCounterForA < mA) {
    // 目标：设置C的第colCounterForB列第rowCounterForA行的元素
    var i = AcolPtrs(rowCounterForA)
    val indEnd = AcolPtrs(rowCounterForA + 1)
    // i和indEnd表示A的第rowCounterForA行有多少元素，
    // 稀疏矩阵A的其他行是0，所以就可以不管了
    var sum = 0.0
    while (i < indEnd) {
      // 因为B是Dense的，这里要算一下实际的坐标，
      // 我们要知道B的是第colCounterForB列，对应到就是A的第colCounterForB行，
      // 所以查询ArowIndices表就能知道对应的行便宜
      sum += Avals(i) * Bvals(Bstart + ArowIndices(i))
      i += 1
    }
    
    val Cindex = Cstart + rowCounterForA
    Cvals(Cindex) = beta * Cvals(Cindex) + sum * alpha
    rowCounterForA += 1
  }
  colCounterForB += 1
}

有gemm就有gemv，也就是C变成Vector了。除此之外，还有对于各种不同m的特化实现。例如SYMV是对称矩阵乘法，TRMV是三角矩阵乘法。com.github.fommil.netlib.BLAS里面有dgbmv带状矩阵，dsbmv，dtrmv等。此外，这些在Breeze库里面也有提供。

启发

  1. transpose可以做到lazy，但得是对二维特化的
  2. 矩阵可以Sparse存储
  3. 矩阵和多维数组最好单独实现

Breeze库

breeze库是一个类似numpy的线性代数库
其中BDM是一个Matrix，BDV是一个Vector，其实都是别名import breeze.linalg.{DenseMatrix => BDM}

其他代码

Attribute对象

Attribute继承自AttributeFactory这个object（原来伴生对象也能继承），主要是从StructField生成Attribute对象。StructField是构成StructType的成员。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
def fromStructField(field: StructField): Attribute = decodeStructField(field, false)

// 下面这个函数主要就是从filed.metadata中解析得到Attribute对象
private[ml] def decodeStructField(field: StructField, preserveName: Boolean): Attribute = {
  require(field.dataType.isInstanceOf[NumericType])
  val metadata = field.metadata
  val mlAttr = AttributeKeys.ML_ATTR // 值为ML_ATTR
  if (metadata.contains(mlAttr)) {
    // 是使用metadata里面的名字，还是field.name的名字
    val attr = fromMetadata(metadata.getMetadata(mlAttr))
    if (preserveName) {
      attr
    } else {
      // Copy with a new name
      attr.withName(field.name)
    }
  } else {
    UnresolvedAttribute
  }
}

private[attribute] override def fromMetadata(metadata: Metadata): Attribute = {
  import org.apache.spark.ml.attribute.AttributeKeys._
  val attrType = if (metadata.contains(TYPE)) {
    metadata.getString(TYPE)
  } else {
    AttributeType.Numeric.name
  }
  // 首先得到一个AttributeFactory，然后结合metadata得到一个Attribute
  getFactory(attrType).fromMetadata(metadata)
}

private def getFactory(attrType: String): AttributeFactory = {
  if (attrType == AttributeType.Numeric.name) {
    NumericAttribute
  } else if (attrType == AttributeType.Nominal.name) {
    NominalAttribute
  } else if (attrType == AttributeType.Binary.name) {
    BinaryAttribute
  } else {
    throw new IllegalArgumentException(s"Cannot recognize type $attrType.")
  }
}

Reference