Skip to content

[SPARK-3097][MLlib] Word2Vec performance improvement #1932

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Ishiihara wants to merge 5 commits into from
Closed

[SPARK-3097][MLlib] Word2Vec performance improvement #1932

Changes from 2 commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
aa2ab36
use reduceByKey to combine models
Ishiihara Aug 13, 2014
9075e1c
combine syn0Global and syn1Global to synGlobal
Ishiihara Aug 13, 2014
083aa66
update synGlobal in place and reduce synOut size
Ishiihara Aug 14, 2014
cad2011
bug fix for synModify array out of bound
Ishiihara Aug 14, 2014
d5377a9
use syn0Global and syn1Global to represent model
Ishiihara Aug 17, 2014
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
61 changes: 32 additions & 29 deletions mllib/src/main/scala/org/apache/spark/mllib/feature/Word2Vec.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ import org.apache.spark.mllib.rdd.RDDFunctions._
import org.apache.spark.rdd._
import org.apache.spark.util.Utils
import org.apache.spark.util.random.XORShiftRandom

import org.apache.spark.util.collection.PrimitiveKeyOpenHashMap
/**

@mengxr mengxr Aug 14, 2014

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

add an empty line after imports

* Entry in vocabulary
*/
Expand Down Expand Up @@ -235,7 +235,7 @@ class Word2Vec extends Serializable with Logging {
b = 0
while (b < i) {
vocab(a).code(i - b - 1) = code(b)
vocab(a).point(i - b) = point(b) - vocabSize
vocab(a).point(i - b) = point(b)
b += 1
}
a += 1
Expand Down Expand Up @@ -284,16 +284,15 @@ class Word2Vec extends Serializable with Logging {

val newSentences = sentences.repartition(numPartitions).cache()
val initRandom = new XORShiftRandom(seed)
var syn0Global =
Array.fill[Float](vocabSize * vectorSize)((initRandom.nextFloat() - 0.5f) / vectorSize)
var syn1Global = new Array[Float](vocabSize * vectorSize)

var synGlobal =

@mengxr mengxr Aug 14, 2014

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we want to keep syn0 and syn1 in order to have an easy mapping from/to the original C implementation? It reduces the code maintenance cost.

@Ishiihara Ishiihara Aug 14, 2014

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We can keep syn0 and syn1, but it adds some unnecessary slicing operations on array.

@mengxr mengxr Aug 14, 2014

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There is no slicing across both syn0 and syn1, right?

@Ishiihara Ishiihara Aug 14, 2014

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We need to perform reduceByKey on both syn0 and syn1 and we have different updated keys for syn0 and syn1. To perform reduceByKey of syn0 and syn1 together, we need to have a unique key and one way to achieve this is to treat i + vocabSize as the key for syn1(i). Then after we collect, we need to slice to update syn0Global and syn1Global. Any better idea?

@mengxr mengxr Aug 14, 2014

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Making composite key in the output RDD so reduceByKey can distinguish whether the update is for syn0 or syn1?

Array.fill[Float](2 * vocabSize * vectorSize)((initRandom.nextFloat() - 0.5f) / vectorSize)
var alpha = startingAlpha
for (k <- 1 to numIterations) {
val partial = newSentences.mapPartitionsWithIndex { case (idx, iter) =>
val random = new XORShiftRandom(seed ^ ((idx + 1) << 16) ^ ((-k - 1) << 8))
val model = iter.foldLeft((syn0Global, syn1Global, 0, 0)) {
case ((syn0, syn1, lastWordCount, wordCount), sentence) =>
val synModify = new Array[Int](2 * vocabSize)
val model = iter.foldLeft((synGlobal, 0, 0)) {
case ((syn, lastWordCount, wordCount), sentence) =>
var lwc = lastWordCount
var wc = wordCount
if (wordCount - lastWordCount > 10000) {
Expand Down Expand Up @@ -321,42 +320,46 @@ class Word2Vec extends Serializable with Logging {
// Hierarchical softmax
var d = 0
while (d < bcVocab.value(word).codeLen) {
val l2 = bcVocab.value(word).point(d) * vectorSize
val ind = bcVocab.value(word).point(d)
val l2 = ind * vectorSize
// Propagate hidden -> output
var f = blas.sdot(vectorSize, syn0, l1, 1, syn1, l2, 1)
synModify(ind) += 1
var f = blas.sdot(vectorSize, syn, l1, 1, syn, l2, 1)
if (f > -MAX_EXP && f < MAX_EXP) {
val ind = ((f + MAX_EXP) * (EXP_TABLE_SIZE / MAX_EXP / 2.0)).toInt
f = expTable.value(ind)
val g = ((1 - bcVocab.value(word).code(d) - f) * alpha).toFloat
blas.saxpy(vectorSize, g, syn1, l2, 1, neu1e, 0, 1)
blas.saxpy(vectorSize, g, syn0, l1, 1, syn1, l2, 1)
blas.saxpy(vectorSize, g, syn, l2, 1, neu1e, 0, 1)
blas.saxpy(vectorSize, g, syn, l1, 1, syn, l2, 1)
}
d += 1
}
blas.saxpy(vectorSize, 1.0f, neu1e, 0, 1, syn0, l1, 1)
blas.saxpy(vectorSize, 1.0f, neu1e, 0, 1, syn, l1, 1)
synModify(lastWord) += 1
}
}
a += 1
}
pos += 1
}
(syn0, syn1, lwc, wc)
(syn, lwc, wc)
}
Iterator(model)
}
val (aggSyn0, aggSyn1, _, _) =
partial.treeReduce { case ((syn0_1, syn1_1, lwc_1, wc_1), (syn0_2, syn1_2, lwc_2, wc_2)) =>
val n = syn0_1.length
val weight1 = 1.0f * wc_1 / (wc_1 + wc_2)
val weight2 = 1.0f * wc_2 / (wc_1 + wc_2)
blas.sscal(n, weight1, syn0_1, 1)
blas.sscal(n, weight1, syn1_1, 1)
blas.saxpy(n, weight2, syn0_2, 1, syn0_1, 1)
blas.saxpy(n, weight2, syn1_2, 1, syn1_1, 1)
(syn0_1, syn1_1, lwc_1 + lwc_2, wc_1 + wc_2)
val synLocal = model._1
val synOut = new PrimitiveKeyOpenHashMap[Int, Array[Float]](vocabSize * 2)
var index = 0
while(index < 2 * vocabSize) {
if (synModify(index) != 0) {
synOut.update(index, synLocal.slice(index * vectorSize, (index + 1) * vectorSize))
}
index += 1
}
syn0Global = aggSyn0
syn1Global = aggSyn1
Iterator(synOut)
}
synGlobal = partial.flatMap(x => x).reduceByKey {
case (v1,v2) =>

@mengxr mengxr Aug 14, 2014

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

move case (v1, v2) => to previous line and add a space after ,

blas.saxpy(vectorSize, 1.0f, v2, 1, v1, 1)
v1
}.collect().sortBy(_._1).flatMap(x => x._2)

@mengxr mengxr Aug 14, 2014

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

updating synGlobal in-place is more memory-efficient. We don't need to allocate new storage and sort.

}
newSentences.unpersist()

Expand All @@ -365,7 +368,7 @@ class Word2Vec extends Serializable with Logging {
while (i < vocabSize) {
val word = bcVocab.value(i).word
val vector = new Array[Float](vectorSize)
Array.copy(syn0Global, i * vectorSize, vector, 0, vectorSize)
Array.copy(synGlobal, i * vectorSize, vector, 0, vectorSize)
word2VecMap += word -> vector
i += 1
}
Expand Down