基於spark mllib的LDA模型訓練源碼解析

一直想寫一篇關於LDA模型訓練的源代碼走讀，但是因爲個人水平以及時間原因未能如願，今天想起來就記錄了一下源碼走讀過程。有什麼解釋的不太清楚或者錯誤的地方請大家指正。

LDA模型訓練大致經過以下這些步驟：

1. 輸入數據（已轉換爲Vector）和參數設置
2. 根據LDA選擇的算法初始化優化器
3. 迭代優化器
4. 獲得LDA模型

下面對每一步的源碼進行代碼跟進。完整的項目可以到我的github下載

 

1. 輸入數據和參數設置

文件：ckooc-ml/algorithm/utils/LDAUtils.scala

入口方法：train()

  

主要是紅框中的三個部分：數據向量化、LDA優化器參數設置、執行訓練

 

這裏document中的Long類型是每個文檔的ID，後面Vector是tokens的向量表示，主要形式是（詞彙表大小，token的index數組，token對應的WC的數組）

 

LDA優化器參數設置主要是對訓練時需要用到的主題數、迭代次數、初始alpha、初始beta等進行設置

 

run方法時LDA訓練的主入口，方法具體實現如下：

主要實現三個功能：優化器的初始化、迭代優化器、獲取模型。後面對這幾部分進行詳細解析。

 

2. 根據LDA選擇的算法初始化優化器

因爲這裏我使用的是EM算法，所以跟進LDAOptimizer.scala中直接看EMLDAOptimizer的initialize(docs: RDD[(Long, Vector)], lda:LDA)方法即可。

詳細解析：

2.1設置參數

 

2.2設置alpha和beta

注意：默認的alpha= (50.0 / K) + 1.0，beta = 1.1.通常情況下不用對這兩個超參數的初始值進行特殊設置，直接使用默認值即可。

 

2.3因爲LDA模型訓練使用的是圖計算，故在此生成圖的邊（Document -> Term）

每條邊包含文檔ID，詞的索引、詞對應的WC，其中term2index方法功能如下：

 

2.4生成圖的各個節點

從上述代碼可以看到每個節點都是由一個節點ID和對應的由隨機函數產生的關於主題的隨機向量組成，節點ID又和邊關聯（VertexId= edge.srcId或edge.dstId）

 

2.5構建圖以及優化器參數設置

 

 

3. 迭代優化器

優化器的迭代主要是由優化器的next()方法實現：

這一步的實現主要使用了EM算法總的來說分爲兩步E-Step和M-Step，這兩步的解釋如下：

• E-Step：假定參數已知，計算此時隱變量的後驗概率。

• M-Step：帶入隱變量的後驗概率，最大化樣本分佈的對數似然函數，求解相應的參數和。

 

詳細解析：

3.1 E-Step，計算每篇文檔的後驗概率，形成一個後驗主題概率分佈

這裏主要計算每個主題的後驗概率的方法時computePTopic()方法：

 

3.2 M-Step，根據後驗概率分佈計算和

 

3.3根據新的和來更新圖，爲下一次迭代做準備

 

 

EM算法的實現也可以參考：“通俗理解LDA主題模型”中關於pLSA和LDA的參數估計部分

 

 

4. 獲得LDA模型

這一部分比較簡單，直接使用訓練好的graph等信息生成一個DistributedLDAModel即可

 

 

5. 總結

以上就是spark上整個的LDA模型訓練過程。總結下來還是挺簡單的，主要就以下幾個步驟：

• 初始參數構造圖（邊：doc-term,頂點：doc-topics）
• 使用EM算法，計算每個doc的，形成doc 的後驗主題分佈
  
• 根據後驗主題分佈計算參數和
• 根據參數和更新圖
• 得到模型

其中2-4步根據迭代次數進行迭代