使用Ray的知識圖譜分布式模型研究

方志聰，李廷元

（中國民用航空飛行學院計算機學院，德陽 618300）

0 引言

知識圖譜（knowledge graphs，KG）已成為當前的一個顛覆性的數據處理方式，提供統一和語義豐富的、龐大和多元化的數據源訪問。大型圖譜的“理解”能力至關重要，在數據集成、實體解析、搜索、數據提取、數據交換等方面中可以體現，表示學習（RL）/KG 已經成為一個關鍵性促進因素。

強化學習［1］的目標是將一個高維但非常稀疏的KG 嵌入到低維空間中，同時保留原始圖的特征。KG 嵌入模型，如TransE［2］、DistMult［3］和ComplEx［4］處理三元組。三元組（s,o,r）代表一個事實，其中s是主語，o是賓語，r是它們之間的關系。這些方法獨立處理知識圖中的每個三元組，它們學習各種結構關系模式，例如實體之間的對稱關系。知識圖表示學習［5］的最新進展表明，基于圖神經網絡的知識圖嵌入方法表現出優于傳統嵌入方法的性能［6］。用于鏈路預測的最先進的KG 嵌入方法［7-8］通常遵循編碼器-解碼器架構，其中編碼器聚合多處上下文以生成節點嵌入，而解碼器學習節點之間的關系。

1 分布式知識圖譜

由于KG 的規模和計算復雜性不斷增長，分布式KG 嵌入訓練最近在研究界引起了相當大的關注?，F有框架［9-11］采用分布式、數據并行架構來應對高計算量和大內存需求。在這些架構中，首先對KG 進行分區，然后將分區分布在計算節點上進行訓練。這些框架采用不同的方法進行訓練，例如使用環AllReduce 通信架構進行梯度交換［10］。PBG［9］和DGL-KE［11］是以傳統的KG嵌入模型設計的，例如TransE 和DistMult，它們獨立處理三元組。這些模型不能用于需要k-hop鄰域信息的基于圖神經網絡的KG 嵌入模型的分布式訓練。圖神經網絡架構分布式訓練的技術挑戰之一是有效地將圖劃分為足夠多的分區，以使每個計算節點都有相對少量的、自給自足的數據可供使用，并且工作負載在所有計算節點之間保持平衡，以避免SSGD期間的空閑時間。

Sheikh 等［10］提出了一種基于圖神經網絡的KG 嵌入模型的分布式訓練模型，使用鄰域展開方法使分區自給自足。首先使用邊緣劃分方法對KG 進行分區，然后使用鄰域展開過程對分區進行擴展。鄰域擴展過程包括了一個分區中所有節點的所有k-hop 鄰居，以避免訓練過程中的交叉分區通信開銷，以節點和邊緣復制為代價。本文還引入了基于約束的負抽樣方法進行訓練。實驗表明，從分區內抽取負樣本是有效的，避免了分區中的負樣本，進一步減少了交叉分區的傳輸。

文獻［12］中的分布式訓練方法基于PyTorch分布式API。使用PyTorch 等底層框架的核心級API，對擴展和部署具有挑戰性。最近的分布式框架，如Ray［13］簡化了擴展和部署的底層細節。Ray是一個通用集群計算框架，提供通用、簡單的API，使開發人員能夠使用現有的庫和系統進行分布式計算。豐富的Ray框架提供了用于訓練（Ray Train）、調優（Ray Tune）、在線服務（Ray Serve）和分布式數據管理（Modin）［14］的API。Ray Train 提供了一個簡單易用的API，用于使用深度學習框架（如PyTorch）對神經網絡模型進行分布式訓練。此外，它提供了可組合性，允許Ray Train與Ray Tune 互操作以進行超參數學習。本文的貢獻總結如下：

（1）介紹了一個使用Ray 的分布式KG 嵌入訓練框架，并在KG上展示了性能結果。

（2）使用Ray提供和評估分布式屬性預處理。

（3）使用RayTrain API 對KG 嵌入模型進行分布式訓練，并在執行復雜任務時展示API 的簡單性。

2 模型概覽

使用Ray 對KG 嵌入模型進行分布式訓練的模型架構如圖1所示。模型的輸入是原始圖、節點屬性和關系以及圖的分區。該模型由兩個主要組件組成：分布式數據預處理和分布式訓練組件。分布式數據預處理組件接收節點的原始屬性，并使用Ray 在分布式設置中處理它們以生成特征向量。分布式訓練組件的目標是學習KG嵌入模型。分布式設置的組件如下所述。

圖1 使用Ray訓練知識圖譜架構

2.1 分布式數據預處理

本文的輸入數據由編碼為特征向量的節點屬性組成，以便訓練圖神經網絡。 這些節點屬性可以是不同的格式，例如文本、數字、日期和分類類型。具體來說，將日期轉換為數值，本文使用scikit-learn KBinsDiscretizer 和One-HotEncoder（或MultiLabelBinarizer）對數值和分類值進行編碼，并使用預訓練的BERT 語言模型對文本屬性進行編碼。

對節點屬性進行編碼的過程分為兩個階段：在第一階段，本文加載了一個預訓練的BERT模型并將KBinsDiscretizer、OneHotEncoder 和MultiLabelBinarizer 模型擬合到輸入數據中的所有節點。學習模型的過程不是以分布式方式執行的。 在第二階段，本文將學習到的模型應用于節點屬性，將它們轉換為特征向量。這個轉換步驟是使用Ray以分布式方式執行的。

本文使用@ray.remote 實現了一個Ray 遠程工作，它接收作為輸入的節點子集及其屬性和對預訓練模型的引用。通過拆分節點數據并將數據分布在多個Ray 通道之間，可以并行化屬性轉換。由于所有Ray 通道使用相同的預訓練scikit-learn 和BERT 模型，因此使用ray.pu（t）將模型放置在分布式對象存儲中，并將返回的對象ID傳遞給通道。

轉換結果是使用ray.get（）從Ray 通道收集的，并連接到一個數組中，其中包含所有數據的轉換節點屬性。

應對實體屬性的異質性，本文將屬性建模為實體的直接鄰居，并使用類似沒有自環的1-hop GCN［15］網絡來獲得實體屬性嵌入。也就是說，本文使用一個虛擬節點來表示一個實體，其屬性嵌入是其屬性鄰居嵌入聚合的結果。與GCN不同，每個屬性類型都有一個專用的嵌入矩陣用于轉換，創建的實體屬性圖僅用于屬性嵌入，在知識嵌入模型訓練期間，實體嵌入用作初始節點特征。

2.2 分布式訓練

本文使用Ray 訓練API，對本文的圖神經網絡模型進行分布式訓練。該API 支持Torch、Tensorflow 和Horovod 庫，使用文獻［12］中描述的訓練架構，并將其封裝在Ray Train中進行分布式訓練，如圖1所示。當Ray通道運行trainier.run（），每個通道加載其圖形分區并初始化模型以進行培訓。通道的數量是由圖分區的數量決定的。使用文獻［12］跟蹤訓練函數邏輯，在每個epoch之后，Ray訓練允許從分布式通道內收集中間結果。本文使用PyTorch 作為后端。因此，Hence,Ray分布式會將梯度共享和更新。

3 實驗環境

本文使用Ray 1.8.0［8］作為分布式訓練框架，PyTorch Geometric 1.7.2［15］作為圖嵌入框架，Py-Torch 1.9.0［12］作為深度學習后端。本文在一個由2臺機器組成的集群上進行了實驗。集群安裝并運行的是CentOS Linux 7.9，并且每臺機器上有一條16 GB RAM 的內存條和一塊RTX1080 GPU圖形處理器。

4 實驗和評估

本文針對鏈路預測任務的非分布式訓練評估了使用Ray 的分布式訓練的準確性和可擴展性性能，在兩種設置中使用了相同的超參數。RelGNN 將鏈接預測視為二元分類任務，即預測給定的三元組是否有效。

在本節中，本文使用第2節中描述的數據描述KG 嵌入模型訓練的評估設置。數據由2.8 KB節點、24 KB 邊和5 種關系類型組成。通過處理節點屬性獲得的每個節點的特征向量的維度為32。為了訓練KG 嵌入模型，本文將數據分為訓練集、測試集和驗證集。本文隨機選擇80%的數據作為訓練集，測試集和驗證集各使用10%的數據。

本文的輸入數據包含1萬多個節點，這些節點均勻分布在Ray 通道之間進行預處理，其中每個通道根據任務分配到一個GPU 或CPU。由于scikit-learn 不支持GPU 計算，因此屬性轉換是在CPU 上計算的，而BERT 轉換是在GPU 上計算的。本文測量了1、2 通道的屬性處理時間。如圖2、圖3 所示，使用1 個通道處理所有節點屬性的進程處理時間和一個epoch 時間分別為104秒和15秒，使用2個通道進程處理時間和一個epoch時間分別減少到35秒和5秒。通道從1增加到2時，處理時間縮短為原來的1/3。

圖2 Ray分布式節點屬性處理時間

圖3 Ray中RelGNN訓練的epoch

使用Ray Trainer API 引入的開銷可以忽略不計，大約在開始訓練的部分時間會有差異。我們還使用pytorch.distributed API 進行了評估，發現這兩種方法的運行時間沒有顯著差異。因此，使用Ray API 開發應用程序可以簡化分布式訓練過程，而無需或可忽略額外的時間成本。

本文評估的訓練精度測試集使用曲線下的面積（AUC）。由于模型和輸入數據都很大，GPU內存有限，所以不能利用GPU 加速。因此，在小區采樣CPU 上執行訓練和使用。如圖4 所示，與非分布式訓練相比，分布式訓練精度幾乎沒有下降。

圖4 Ray中RelGNN訓練的AUC分數

5 結語

本文實現了一個分布式模型，使用Ray API對知識嵌入模型進行分布式訓練。 Ray 提供了易于使用、功能強大且靈活的API，它們隱藏了用于擴展和部署的底層細節。本文有效地使用這些API 進行分布式數據預處理和訓練。對Ray集群的實驗評估表明，本文能夠在數據預處理任務上實現1 倍的加速。在分布式訓練任務上，本文在Rel 圖神經網絡上實現3 倍的加速，并獲得了與非分布式訓練設置相同的分數。Ray API引入的開銷可以忽略不計。在接下來的工作中，計劃利用Ray 支持對輸入知識圖譜的可擴展模型推理，進一步對知識圖譜分布式訓練進行加速。