方志聰，李廷元

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

0 引言

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

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

1 分布式知識圖譜

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

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

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

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

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

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

2 模型概覽

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

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

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

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

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

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

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

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

2.2 分布式訓練

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

3 實驗環(huán)境

本文使用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 的內(nèi)存條和一塊RTX1080 GPU圖形處理器。

4 實驗和評估

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

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

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

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

圖3 Ray中RelGNN訓練的epoch

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

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

圖4 Ray中RelGNN訓練的AUC分數(shù)

5 結(jié)語

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