北大、微軟亞洲研究院：高效的大規模圖神經網路計算

新聞 11-01

新智元報道

來源：arXiv

編輯：肖琴

【新智元導讀】第一個支持圖神經網路的並行處理框架出現了！北京大學、微軟亞洲研究院的研究人員近日發表論文，提出NGra，這是第一個支持大規模GNN的系統。

GNN（圖神經網路）代表了一種新興的計算模型，這自然地產生了對在大型graph上應用神經網路模型的需求。

但是，由於GNN固有的複雜性，這些模型超出了現有深度學習框架的設計範圍。此外，這些模型不容易在並行硬體（如GPU）上有效地加速。

近日，北京大學、微軟亞洲研究院的多位研究人員在arXiv上發布了一篇新論文，提出了解決這些問題的有效方案。

論文題為Towards Efficient Large-Scale Graph Neural Network Computing：

論文地址：https://arxiv.org/pdf/1810.08403.pdf

作者表示：「我們提出NGra，這是第一個基於圖形的深度神經網路並行處理框架。」

NGra描述了一種新的SAGA-NN模型，用於將深度神經網路表示為頂點程序（vertex programs），其中每一層都在明確定義的圖形操作階段（Scatter，ApplyEdge，Gather，ApplyVertex）。

這個模型不僅允許直觀地表示GNN，而且還可以方便地映射到高效的數據流表示。NGra通過GPU核心或多GPU的自動圖分區和基於chunk的流處理透明地解決了可擴展性挑戰，仔細考慮了數據局部性、數據移動以及並行處理和數據移動的重疊。

NGra通過在GPU上進行高度優化的Scatter / Gather操作進一步提高了效率，儘管它具有稀疏性。我們的評估表明，NGra可以擴展到現有框架無法直接處理的大型實際圖形，而在TensorFlow的multiple-baseline設計上，即使在小規模上也可以實現約4倍的加速。

第一個支持大規模GNN的系統

NGra是第一個支持大規模GNN（圖神經網路）的系統，這是一個在GPU上可擴展、高效的並行處理引擎。

NGra自然地將數據流（dataflow）與頂點程序抽象（vertex-program abstraction）結合在一個新模型中，我們將其命名為SAGA-NN（Scatter-ApplyEdge-Gather-ApplyVertex with Neural Networks）。

雖然SAGA可以被認為是GAS（Gather-Apply-Scatter）模型的變體，但SAGA-NN模型中的用戶定義函數允許用戶通過使用數據流抽象來表示對vertex或edge數據（被視為tensors）的神經網路計算，而不是專為傳統圖形處理而設計（例如PageRank、 connected component和最短路徑等演算法）

與DNN一樣，高效地使用GPU對於GNN的性能至關重要，而且由於要處理的是大型圖形結構，這一點更為重要。為了實現超出GPU物理限制的可擴展性，NGra將圖形（頂點和邊緣數據）透明地劃分為塊（chunk），並將SAGA-NN模型中表示的GNN演算法轉換為具有chunk粒度的運算符的dataflow graph，從而在單個GPU或多個GPU上啟用基於chunk的並行流處理。

NGra engine的效率在很大程度上取決於NGra如何管理和調度並行流處理，以及在GPU上關鍵圖形傳播運算符Scatter和Gather的實現。

NGra非常注重數據局部性，以最大限度地減少GPU內存中的數據交換，並在GPU內存中最大化數據塊的重用，同時將數據移動和計算以流的方式重疊。

對於多GPU的情況，它使用 ring-based streaming機制，通過直接在GPU之間交換數據塊來避免主機內存中的冗餘數據移動。

與其他基於GPU的圖形引擎關注的傳統圖形處理場景不同，在GNN場景中，可變頂點數據本身可能無法容納到GPU設備內存中，因為每個頂點的數據可以是特徵向量（ feature vector）而不是簡單的標量（scalar）。因此，我們的方案更傾向於在每個頂點數據訪問中利用並行性，從而提高內存訪問效率。

我們通過使用vertex-program abstraction和圖形傳播過程的自定義運算符擴展TensorFlow，從而實現NGra。

我們利用單個伺服器的主機內存和GPU的計算能力，證明NGra可以擴展以支持大型圖形的各種GNN演算法，其中許多是現有深度學習框架無法直接實現的。

與小型graph上的TensorFlow相比，它可以支持GPU，NGra可以獲得最多4倍的加速。我們還廣泛評估了NGra的多重優化所帶來的改進，以證明其有效性。

接下來的部分將描述SAGA-NN編程抽象，NGra系統的組件，以及NGra的實現和評估。

NGra程序抽象

基於圖（graph）的神經網路（GNN）是根據圖形結構定義的一類通用神經網路架構。

圖中的每個頂點或邊可以與張量數據（通常是vector）相關聯，作為其特徵或嵌入。GNN可以堆疊在多個層中，迭代傳播過程在同一個圖上逐層進行。

在圖的每個層中，頂點或邊緣要素沿邊緣變換和傳播，並在目標頂點聚合，以生成下一層的新要素。轉換可以是任意的DNN計算。

圖還可以包含每個頂點，每個邊緣或整個圖形的標籤，用於計算頂層的損失函數。然後從底層到頂層執行前饋計算（feedforward computation）和反向傳播。

圖1描述了一個2層的GNN的前饋計算。

圖1

我們使用Gated Graph ConvNet（G-GCN）演算法作為一個具體示例。 Graph ConvNet概括了卷積運算的概念，通常應用於圖像數據集，用於處理任意圖形（例如knowledge graph）。Gated Graph ConvNet進一步結合了門控機制，因此模型可以了解哪些邊對學習目標更重要。

G-GCN每一層的前饋計算如圖2所示：

圖2：SAGA-NN模型中，Gated Graph ConvNet的layer，其中?指矩陣乘法。

NGra系統的組成

NGra提供dataflow和vertex program abstractions的組合作為用戶界面。

NGra主要包括：

一個前端，它將SAGA-NN模型中實現的演算法轉換為塊粒度數據流圖（chunk-granularity dataflow graph），使GPU中大型圖的GNN計算成為可能；

一個優化層，它產生用於最小化主機和GPU設備存儲器之間的數據移動的調度策略，並識別融合操作和刪除冗餘計算；

一組有效的傳播操作內核，支持基於流的處理，以將GPU中的數據移動和計算重疊；

dataflow execution runtime。NGra主要利用現有的基於數據流的深度學習框架來處理dataflow execution runtime。

圖3：SAGA-NN Stages for each layer of GN

NGra的優化

圖4描述了ApplyEdge階段中矩陣乘法運算：

圖4

圖5顯示了優化的dataflow graph，其中矩陣乘法移入ApplyVertex stage：

圖

圖7是多GPU的架構

圖7：多GPU架構

NGra的評估

我們在TensorFlow (v1.7) 上實現NGra，使用大約2,900行C++代碼和3000行Python代碼。NGra通過前端擴展TensorFlow，將SAGA-NN程序轉換為chunk-granularity dataflow graph，幾個scatter/gather 運算符，以實現高效的圖傳播，以及ring-based的流調度方案。

以下是評估結果。評估證明了NGra的高效和可擴展性，以及與state-of-the-art的系統TensorFlow的比較。