【Spark原理系列】Spark Word2Vec原理示例源码分析详解

【Spark原理系列】$Spark$ $Word2Vec$原理参数示例源码分析详解 点击这里看全文

$Word2Vec$

$Word2Vec$ 是一种计算单词的分布式向量表示的方法。分布式表示的主要优势在于相似的单词在向量空间中靠近，这使得对新模式的泛化更容易，模型估计更稳健。分布式向量表示已被证明在许多自然语言处理应用中非常有用，如命名实体识别、消歧、解析、标记和机器翻译。

模型

在我们的$Word2Vec$实现中，我们使用了$skip-gram$ 模型。$skip-gram$ 的训练目标是学习单词向量表示，在同一个句子中预测其上下文。数学上，给定一系列训练单词 $w1,w2,\dots ,wT$ ，$skip-gram$ 模型的目标是最大化平均对数似然函数

$$\frac{1}{T}\sum _{t=1}^T\sum _{j=-k}^{j=k}logp(\frac{w_{t+j}}{w_t})$$
其中 $k$ 是训练窗口的大小。

在 $skip-gram$ 模型中，每个单词 $w$ 与两个向量 $u_w$ 和 $v_w$ 相关联，它们分别是 $w$ 的单词和上下文的向量表示。正确预测给定单词 $w_i$ 的概率取决于 $softmax$ 模型，即

$$p(wi|wj)=\frac{\exp (u_{w_i}^T\cdot v_{w_j})}{{\sum }_{l=1}^{V}exp(u_l^{\top }{v}_{w_j})}$$

其中$V$ 是词汇表大小。

$skip-gram$ 模型使用$softmax$开销较大，因为计算 $\log p(wi|wj)$ 的成本与 $V$ 成正比，而 $V$ 可能很容易达到数百万级别。为了加速$Word2Vec$的训练，我们使用了分层softmax，将计算$\log p(wi|wj)$的复杂度降低到 $O(log(V))$。

参数方法

Word2VecBase是一个trait，定义了Word2Vec和Word2VecModel的参数。

下面是它的主要方法：

• vectorSize: 词向量的维度，默认为100。
• windowSize: 窗口大小（上下文词的范围），默认为5。
• numPartitions: 分区数，默认为1。
• minCount: 词汇表中一个词必须出现的最少次数，默认为5。
• maxSentenceLength: 输入数据中每个句子的最大长度（以单词计），超过该阈值的句子将被分割成多个片段，默认为1000。

Word2Vec是一个Estimator，用于训练一个将单词转换为向量表示的模型。

以下是它的一些重要方法：

• setInputCol、setOutputCol、setVectorSize等：设置参数的方法。
• fit(dataset: Dataset[_]): Word2VecModel：使用给定的数据集进行训练，并返回训练得到的Word2VecModel。
• transformSchema(schema: StructType): StructType：根据输入模式转换并返回输出模式。

Word2VecModel是由Word2Vec训练得到的模型，它包含一个Map(String, Vector)，将单词映射为向量表示。

Word2Vec和Word2VecModel都提供了序列化和反序列化的方法。

Word2Vec的伴生对象提供了加载模型的方法。

示例

下面的示例演示了如何加载文本文件，将其解析为 Seq[String] 类型的 RDD，构建 Word2Vec 实例，然后用输入数据拟合一个 Word2VecModel。最后，我们展示指定单词的前40个同义词。要运行这个示例，首先下载 text8 数据并将其解压到您喜欢的目录中。这里我们假设提取的文件是 text8，并且与您运行 Spark shell 的目录相同。

数据sample_lda_data.txt

1 2 6 0 2 3 1 1 0 0 3
1 3 0 1 3 0 0 2 0 0 1
1 4 1 0 0 4 9 0 1 2 0
2 1 0 3 0 0 5 0 2 3 9
3 1 1 9 3 0 2 0 0 1 3
4 2 0 3 4 5 1 1 1 4 0
2 1 0 3 0 0 5 0 2 2 9
1 1 1 9 2 1 2 0 0 1 3
4 4 0 3 4 2 1 3 0 0 0
2 8 2 0 3 0 2 0 2 7 2
1 1 1 9 0 2 2 0 0 3 3
4 1 0 0 4 5 1 3 0 1 0

示例代码

import org.apache.spark.mllib.feature.{
   Word2Vec, Word2VecModel}

val input = sc.textFile("data/mllib/sample_lda_data.txt").map(line => line.split(" ").toSeq)

val word2vec = new Word2Vec()

val model = word2vec.fit(input)

val synonyms = model.findSynonyms("1", 5)

for((synonym, cosineSimilarity) <- synonyms) {
   
  println(s"$synonym $cosineSimilarity")
}

// Save and load model
model.save(sc, "myModelPath")
val sameModel = Word2VecModel.load(sc, "myModelPath")

源码

注意：ml中的word2vec源码训练模型的时候调用复用了mllib中的word2vec

ml中Word2Vec源码

package org.apache.spark.ml.feature

import org.apache.hadoop.fs.Path

import org.apache.spark.annotation.Since
import org.apache.spark.internal.config.Kryo.KRYO_SERIALIZER_MAX_BUFFER_SIZE
import org.apache.spark.ml.{
   Estimator, Model}
import org.apache.spark.ml.linalg.{
   BLAS, Vector, Vectors, VectorUDT}
import org.apache.spark.ml.param._
import org.apache.spark.ml.param.shared._
import org.apache.spark.ml.util._
import org.apache.spark.mllib.feature
import org.apache.spark.sql.{
   DataFrame, Dataset, SparkSession}
import org.apache.spark.sql.functions._
import org.apache.spark.sql.types._
import org.apache.spark.util.{
   Utils, VersionUtils}
/**
 * Word2Vec和Word2VecModel的参数。
 */
private[feature] trait Word2VecBase extends Params
  with HasInputCol with HasOutputCol with HasMaxIter with HasStepSize with