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

方志聰，李廷元

（中國民用航空飛行學(xué)院計算機(jī)學(xué)院，德陽 618300）

0 引言

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

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

1 分布式知識圖譜

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

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

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

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

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

（3）使用RayTrain API 對KG 嵌入模型進(jìn)行分布式訓(xùn)練，并在執(zhí)行復(fù)雜任務(wù)時展示API 的簡單性。

2 模型概覽

使用Ray 對KG 嵌入模型進(jìn)行分布式訓(xùn)練的模型架構(gòu)如圖1所示。模型的輸入是原始圖、節(jié)點(diǎn)屬性和關(guān)系以及圖的分區(qū)。該模型由兩個主要組件組成：分布式數(shù)據(jù)預(yù)處理和分布式訓(xùn)練組件。分布式數(shù)據(jù)預(yù)處理組件接收節(jié)點(diǎn)的原始屬性，并使用Ray 在分布式設(shè)置中處理它們以生成特征向量。分布式訓(xùn)練組件的目標(biāo)是學(xué)習(xí)KG嵌入模型。分布式設(shè)置的組件如下所述。

圖1 使用Ray訓(xùn)練知識圖譜架構(gòu)

2.1 分布式數(shù)據(jù)預(yù)處理

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

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

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

轉(zhuǎn)換結(jié)果是使用ray.get（）從Ray 通道收集的，并連接到一個數(shù)組中，其中包含所有數(shù)據(jù)的轉(zhuǎn)換節(jié)點(diǎn)屬性。

應(yīng)對實(shí)體屬性的異質(zhì)性，本文將屬性建模為實(shí)體的直接鄰居，并使用類似沒有自環(huán)的1-hop GCN［15］網(wǎng)絡(luò)來獲得實(shí)體屬性嵌入。也就是說，本文使用一個虛擬節(jié)點(diǎn)來表示一個實(shí)體，其屬性嵌入是其屬性鄰居嵌入聚合的結(jié)果。與GCN不同，每個屬性類型都有一個專用的嵌入矩陣用于轉(zhuǎn)換，創(chuàng)建的實(shí)體屬性圖僅用于屬性嵌入，在知識嵌入模型訓(xùn)練期間，實(shí)體嵌入用作初始節(jié)點(diǎn)特征。

2.2 分布式訓(xùn)練

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

3 實(shí)驗(yàn)環(huán)境

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

4 實(shí)驗(yàn)和評估

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

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

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

圖2 Ray分布式節(jié)點(diǎn)屬性處理時間

圖3 Ray中RelGNN訓(xùn)練的epoch

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

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

圖4 Ray中RelGNN訓(xùn)練的AUC分?jǐn)?shù)

5 結(jié)語

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