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

DinK發表於2018-11-02

微軟神經網路

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的比較。