Spark中的矩陣乘法分析

標籤：

前言：矩陣乘法在資料採礦/機器學習中是常用的計算步驟，並且在大資料計算中，shuffle過程是不可避免的，矩陣乘法的不同計算方式shuffle的資料量都不相同。通過對矩陣乘法不同計算方式的深入學習，希望能夠對大資料演算法實現的shuffle過程最佳化有所啟發。網上有很多分布式矩陣乘法相關的文章和論文，但是鮮有對Spark中分布式矩陣乘法的分析。本文針對Spark中分布式矩陣乘法的實現進行必要的說明討論。
分布式矩陣乘法原理：矩陣乘法計算可以分為內積法和外積法。根據實現顆粒度的不同，也可以分為普通矩陣實現和分塊矩陣實現。矩陣乘法計算公式：
公式1為內積法計算公式，即使用A矩陣的行和B矩陣的列進行向量內積計算，每次計算矩陣C中的一個元素。公式2為外積法計算公式，即使用A矩陣的列和B矩陣的行進行向量外積計算，每次計算一個n×k的矩陣，再將計算出的m個矩陣相加，得到矩陣C。分布式矩陣乘法內積法實現：內積要求每次能夠計算矩陣C中的一個元素，由於C中每個元素的計算是相互獨立的，所以這個計算過程能夠並發執行，由於C中有n×k個元素，所以能夠支援的最大並發量是n×k。在計算C中每一個元素的時候，都需要用到A中某一行的m個元素和B中某一行的m個元素，這個過程需要從A中shuffle m個元素和B中shuffle m個元素到C的計算節點上，所以使用內積法直接計算矩陣的乘法，最多需要shuffle的元素個數為2×m×n×k個（因為C有n×k個元素），每個計算節點一次最少要計算C中的1個元素，所以每個計算節點記憶體最少要能夠儲存2m的資料。但是在大矩陣的運算中，通常並發度不可能開發n×k的程度，所以事實上需要shuffle的資料會遠小於2×m×n×k個，因為由於並發量遠小於C中的元素個數，所以會在同一個計算節點上計算C的多個元素值，這個時候，這個節點上的資料就能重複使用，不需要為每個元素的計算都shuffle 2m的資料量。這也是使用內積發計算分塊矩陣的時候，shuffle量會大大減少的原因。演算法shuffle過程如下：
這個演算法的缺點顯而易見，shuffle的資料過多。在分散式運算系統中，shuffle過程是影響系統效能的重要因素。但是這種演算法在A或B其中一個矩陣是小矩陣的時候是有明顯的優勢的。例如：如果B是小矩陣，矩陣B可以broadcast到分布式矩陣A的每一個節點上，這麼在計算C矩陣的時候，A不需要再進行shuffle操作，能夠充分的利用資料本地性進行計算。分布式矩陣乘法外積法實現：向量外積計算公式如下：
結合外積的計算公式，從上面公式2可以看出，矩陣乘法的外積公式是先計算m個n×k的矩陣，再將m個矩陣相加得到矩陣C。在這個過程中m個矩陣的計算過程是相互獨立的。所以能夠支援的最大並發量是m。在計算m個矩陣的任意一個的時候，都需要用到A中某一列的n個元素和B中某一行的k個元素，所以在計算m個矩陣的時候，最多需要移動m×(n+k)個元素。在計算完m個矩陣之後，需要將m個矩陣加和。矩陣加法是對應位置元素相加，所以在最後計算C中每一個元素的時候，需要將m個矩陣對應位置的資料shuffle到某一個計算節點進行加和，所以在加和過程中，需要shuffle的最大資料量是m×n×k(因為m個矩陣，每個矩陣有n×k個元素)。每個計算節點最少要計算m個矩陣中的1個，所以計算節點最少需要能夠儲存n+k個元素的記憶體，由於每個節點輸出元素量為n×k，這個資料量在大規模矩陣中是相當大的，單個計算節點記憶體很難儲存下來，所以每個節點計算的輸出結果通常是達到一定量後儲存到磁碟上，但是如果是在大規模疏鬆陣列中，m個矩陣中每個矩陣中有值的個數通常都會遠小於n×k個，所以第二次shuffle的資料量通常遠小於n×k。演算法shuffle過程如下：
這個演算法的缺點就是，在分布式大矩陣中，如果不是疏鬆陣列，計算出來的中間矩陣會非常大，要在單個計算節點完全使用記憶體計算中間矩陣基本不可能，需要使用磁碟輔助儲存中間矩陣的計算結果。演算法的優勢在於計算疏鬆陣列和雖然AB都是大規模矩陣，但是計算結果是個小矩陣的時候，這兩種情況每個中間矩陣都能夠完全儲存在記憶體中，就會比較快。分塊矩陣的分布式乘法實現：分塊矩陣乘法的過程和不分塊的計算過程差不多，也可以使用內積法和外積法實現。使用分塊矩陣實現分布式矩陣乘法的好處主要有兩個，一個是減少shuffle過程的資料量，另一個是分塊矩陣在每個小塊在本地計算的時候能夠調用現有矩陣計算包，成熟的矩陣計算包通常來說計算效率都比自己實現的好，例如矩陣計算的常用包BLAS包。分塊矩陣的基本計算公式如下（來自維基百科： https://zh.wikipedia.org/wiki/%E5%88%86%E5%A1%8A%E7%9F%A9%E9%99%A3）：
以上公式的計算複雜度為O(n^3)，除了以上公式之外1969年Strassen利用分治演算法將分塊矩陣乘法的計算複雜度降低到O(n^log7)，計算公式如下所示（來自維基百科：https://zh.wikipedia.org/wiki/%E6%96%BD%E7%89%B9%E6%8B%89%E6%A3%AE%E6%BC%94%E7%AE%97%E6%B3%95）：
java中的Jampack包和JAMA都是使用Strassen演算法實現矩陣乘法，這些軟體包都沒有實現並行計算，但是在分散式運算中可以將每個小塊分布到本地計算，再調用對應的單機包。分塊矩陣理論上，一個矩陣中的行列數是可以不一樣的，但是為了實現方便，分布式分塊矩陣乘法中的實現通常一個矩陣中每個塊的行列數都是一樣的。在Spark中，Spark內建的org.apache.spark.mllib.linalg.distributed.BlockMatrix實現了分布式矩陣乘法，BlockMatrix是使用內積法實現的分布式分塊矩陣的乘法。除此之外，第三方包實現有，南京大學PASA實驗室在Spark上改進的分布式矩陣乘法(網址：http://pasa-bigdata.nju.edu.cn/project/Marlin.html)，PASA的包也是使用內積發實現的分布式分塊矩陣乘法。分塊矩陣的內積法和外積法演算法資料shuffle的和上面提到的非分塊矩陣的一樣，只是每個元素都改成一個小矩陣，這裡只分析shuffle的資料量。假設分塊矩陣A和B還是原來的矩陣，只是分塊成行數r，列數j的小矩陣。則內積法shuffle的資料量為2×m/c×n/r×k/r×(r×c)=2×m×n×k/r,外積法shuffle資料量為m/c×(n/r+k/r+n/r×k/r)×(r×c) = m×(n+k/r+n×k/r)。可見分塊矩陣乘法的演算法shuffle資料量比不分塊的少。Spark中的分塊矩陣乘法不使用外積法實現，主要考慮到外積法記憶體佔用量大。
Spark內建BlockMatrix乘法源碼分析：必要的注釋已經在源碼中給出def multiply(other: BlockMatrix): BlockMatrix = {
.......
if (colsPerBlock == other.rowsPerBlock) {
//GridPartitioner一共分為numRowBlocks*other.numColBlocks個partition
valresultPartitioner = GridPartitioner(numRowBlocks,other.numColBlocks,
math.max(blocks.partitions.length,other.blocks.partitions.length))
// 這裡是計算每個leftDestinations和rightDestinations的類型都是Map[(Int,Int),Set[Int]],也就是先計算左右矩陣的
// 每一塊會shuffle到哪個partition
val(leftDestinations, rightDestinations) = simulateMultiply(other,resultPartitioner)
// Each block of A must be multiplied with the corresponding blocks in the columns of B.
valflatA = blocks.flatMap { case((blockRowIndex, blockColIndex),block) =>
val destinations = leftDestinations.getOrElse((blockRowIndex,blockColIndex), Set.empty)
destinations.map(j => (j, (blockRowIndex,blockColIndex, block)))
}
// Each block of B must be multiplied with the corresponding blocks in each row of A.
valflatB = other.blocks.flatMap { case((blockRowIndex, blockColIndex),block) =>
val destinations = rightDestinations.getOrElse((blockRowIndex,blockColIndex), Set.empty)
destinations.map(j => (j, (blockRowIndex,blockColIndex, block)))
}
// GridPartitioner一共有numRowBlocks*other.numColBlocks 個分區,所以在cogroup的時候，在計算A*B=C的時候，C矩陣所用到的所有A和B中的
//分塊都會在一個partition中，在reduceByKey的時候就可以進行combineByKey進行最佳化，事實上在reduceByKey的過程中，只有相加的過程，
// 沒有shuffle的過程。
valnewBlocks = flatA.cogroup(flatB, resultPartitioner).flatMap { case (pId,(a, b)) =>
a.flatMap { case (leftRowIndex,leftColIndex, leftBlock) =>
b.filter(_._1 == leftColIndex).map { case (rightRowIndex, rightColIndex,rightBlock) =>
//在進行矩陣乘法實現的時候，本地矩陣計算使用com.github.fommil.netlib包提供的矩陣演算法，矩陣加法調用的是scalanlp包提供的矩陣加法
valC = rightBlock match {
case dense: DenseMatrix => leftBlock.multiply(dense)
case sparse: SparseMatrix => leftBlock.multiply(sparse.toDense)
case _ =>
throw new SparkException(s"Unrecognized matrix type${rightBlock.getClass}.")
}
((leftRowIndex, rightColIndex),C.toBreeze)
}
}
}.reduceByKey(resultPartitioner, (a,b) => a + b).mapValues(Matrices.fromBreeze)
// TODO: Try to use aggregateByKey instead of reduceByKey to get rid of intermediate matrices
newBlockMatrix(newBlocks, rowsPerBlock,other.colsPerBlock, numRows(),other.numCols())
} else {
.......
}
}
以上代碼有一個simulateMultiply方法比較重要，源碼注釋如下：private[distributed] def simulateMultiply(
other: BlockMatrix,
partitioner: GridPartitioner): (BlockDestinations,BlockDestinations) = {
val leftMatrix = blockInfo.keys.collect() // blockInfo should already be cached
valrightMatrix = other.blocks.keys.collect()
//以下這段代碼這樣理解，假設A*B=C,因為A11在計算C11到C1n的時候會用到，所以A11在計算C11到C1n的機器都會存放一份。
valleftDestinations = leftMatrix.map { case(rowIndex, colIndex) =>
//左矩陣中列號會和右矩陣行號相同的塊相乘，得到所有右矩陣中行索引和左矩陣中列索引相同的矩陣的位置。
// 由於有這個判斷，右矩陣中沒有值的快左矩陣就不會重複復制了，避免了零值計算。
valrightCounterparts = rightMatrix.filter(_._1 == colIndex)
// 因為矩陣乘完之後還有相加的操作(reduceByKey),相加的操作如果在同一部機器上可以用combineBy進行最佳化，
// 這裡直接得到每一個分塊在進行完乘法之後會在哪些partition中用到。
valpartitions = rightCounterparts.map(b => partitioner.getPartition((rowIndex,b._2)))
((rowIndex, colIndex),partitions.toSet)
}.toMap
val rightDestinations = rightMatrix.map {case (rowIndex,colIndex) =>
val leftCounterparts = leftMatrix.filter(_._2 == rowIndex)
val partitions = leftCounterparts.map(b => partitioner.getPartition((b._1,colIndex)))
((rowIndex, colIndex),partitions.toSet)
}.toMap
(leftDestinations, rightDestinations)
}

從代碼中可以知道，Spark中內建的分塊矩陣乘法要求每個Executor的記憶體最少能夠存下左矩陣一行中所有非零塊和右矩陣一列中的所有非零塊。在計算的過程中只需要一次shuffle。
PASA基於Spark實現的BlockMatrix乘法：源碼注釋如下：def multiply(other: BlockMatrix): BlockMatrix = {
.......
if (numBlksByCol() == other.numBlksByRow()) {
//num of rows to be split of this matrix
valmSplitNum = numBlksByRow()
//num of columns to be split of this matrix, meanwhile num of rows of that matrix
valkSplitNum = numBlksByCol()
//num of columns to be split of that matrix
valnSplitNum = other.numBlksByCol()
val partitioner = new MatrixMultPartitioner(mSplitNum,kSplitNum, nSplitNum)

val thisEmitBlocks = blocks.flatMap({case(blkId,blk) =>
//左矩陣的每一塊要與右矩陣對應行的每一塊矩陣相乘，右矩陣每一行有nSplitNum個塊，所以每個塊複製nSplitNum份。
//這裡沒有考慮右矩陣某些塊是零值，所以在join進行本地矩陣乘法的時候會有不必要零值計算
Iterator.tabulate[(BlockID,SubMatrix)](nSplitNum)(i => {
val seq = blkId.row * nSplitNum * kSplitNum + i * kSplitNum + blkId.column
(BlockID(blkId.row,i, seq), blk)})
}).partitionBy(partitioner)
val otherEmitBlocks = other.blocks.flatMap({case(blkId,blk) =>
Iterator.tabulate[(BlockID,SubMatrix)](mSplitNum)(i => {
val seq = i * nSplitNum * kSplitNum + blkId.column * kSplitNum + blkId.row
(BlockID(i, blkId.column, seq),blk)
})
}).partitionBy(partitioner)
if (kSplitNum != 1) {
//以下代碼join用到的Partitioner是MatrixMultPartitioner，reduceByKey用到的Partitioner是HashPartitioner，
// 兩次shuffle用到不同的Partitioner不同，所以不可避免的需要兩次shuffle。
valresult = thisEmitBlocks.join(otherEmitBlocks).mapPartitions(iter =>
iter.map { case (blkId,(block1, block2)) =>
(BlockID(blkId.row,blkId.column), block1.multiply(block2))
}
).reduceByKey((a, b) => a.add(b))
new BlockMatrix(result,numRows(), other.numCols(),mSplitNum, nSplitNum)
} else {
val result = thisEmitBlocks.join(otherEmitBlocks).mapPartitions(iter =>
iter.map { case (blkId,(block1, block2)) =>
(BlockID(blkId.row,blkId.column), block1.multiply(block2))
}
)
new BlockMatrix(result,numRows(), other.numCols(),mSplitNum, nSplitNum)
}
}
......
}

總結：    通過源碼分析，可以知道Spark內建的BlockMatrix乘法演算法比PASA實現的BlockMatrix乘法演算法更高效，能夠避免不必要的零值計算，也能夠減少一次shuffle。在實踐中，使用Spark內建的BlockMatrix演算法要注意記憶體的使用，分塊的時候，塊的大小是多少除了注意記憶體之外，還要注意令子塊中的資料能夠盡量的緊湊，減少零值計算。

Spark中的矩陣乘法分析